import logging
import functools
from typing import Optional
from datetime import timedelta, datetime
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse

import requests
from requests.adapters import HTTPAdapter, Retry
from rest_framework.request import Request
from django.urls import reverse
from django.utils import timezone
from django.utils.crypto import get_random_string
from django.utils.http import url_has_allowed_host_and_scheme
from django.http import HttpResponseRedirect, HttpResponse
from django.contrib.auth import REDIRECT_FIELD_NAME
from django.conf import settings

from labs.models import LabsOAuthTiDBCloudGrant, UserServerlessCluster
from labs.social.tidbcloud import TiDBCloudOAuth2
from labs.tidb_serverless.exceptions import (
    TiDBCloudAPIException,
    TiDBCloudUnauthorizedException,
    TiDBCloudRateLimitException,
    TiDBCloudServerlessException,
)

logger = logging.getLogger(__name__)

TIDBCLOUD_OAUTH_STATE_KEY = "tidbcloud_oauth_state"
TIDBCLOUD_OAUTH_NEXT_URL_KEY = "tidbcloud_oauth_next_url"


def get_tidbcloud_oauth_redirect_uri(request: Request):
    redirect_uri = request.build_absolute_uri(reverse("labs-oauth-tidbcloud-complete"))
    return redirect_uri


def tidbcloud_oauth_start(request: Request) -> HttpResponseRedirect:
    next_url = request.GET.get(REDIRECT_FIELD_NAME)
    if next_url:
        request.session[TIDBCLOUD_OAUTH_NEXT_URL_KEY] = next_url
    tidbcloud_oauth = TiDBCloudOAuth2()
    tidbcloud_oauth.DEFAULT_SCOPE = ["org:member", "project:owner"]
    tidbcloud_oauth.redirect_uri = get_tidbcloud_oauth_redirect_uri(request)
    state = tidbcloud_oauth.get_or_create_state()
    request.session[TIDBCLOUD_OAUTH_STATE_KEY] = state
    result = tidbcloud_oauth.start()
    return HttpResponseRedirect(result.url)


def tidbcloud_oauth_complete(request: Request) -> HttpResponse:
    next_url = request.session.get(TIDBCLOUD_OAUTH_NEXT_URL_KEY)
    if not url_has_allowed_host_and_scheme(
        next_url, allowed_hosts=settings.ALLOWED_HOSTS
    ):
        next_url = "/lab"
    next_url_parsed = list(urlparse(next_url))
    next_url_query = dict(parse_qs(next_url_parsed[4]))

    def success_response(grant: LabsOAuthTiDBCloudGrant):
        next_url_query.update({"serverless_grant_id": grant.id})
        next_url_parsed[4] = urlencode(next_url_query, doseq=True)
        next_url = urlunparse(next_url_parsed)
        return HttpResponseRedirect(next_url)

    def error_response(error: str):
        next_url_query.update(
            {
                "serverless_grant_error_code": "oauth_error",
                "serverless_grant_error_message": error,
            }
        )
        next_url_parsed[4] = urlencode(next_url_query, doseq=True)
        next_url = urlunparse(next_url_parsed)
        return HttpResponseRedirect(next_url)

    state = request.query_params.get("state")
    error = request.query_params.get("error")
    if error:
        return error_response(error)

    if state != request.session.get(TIDBCLOUD_OAUTH_STATE_KEY):
        logger.info(
            f"TiDB Cloud OAuth state mismatch: {state} != {request.session.get(TIDBCLOUD_OAUTH_STATE_KEY)}"
        )
        return error_response("OAuth state mismatch")

    oauth = TiDBCloudOAuth2()
    oauth.data = request.query_params
    oauth.redirect_uri = get_tidbcloud_oauth_redirect_uri(request)
    try:
        response = oauth.request_access_token(
            oauth.access_token_url(),
            data=oauth.auth_complete_params(state=state),
            params=None,
            headers=oauth.auth_headers(),
            auth=oauth.auth_complete_credentials(),
            method=oauth.ACCESS_TOKEN_METHOD,
        )
        grant = LabsOAuthTiDBCloudGrant.objects.create(
            user=request.user,
            access_token=response.get("access_token"),
            refresh_token=response.get("refresh_token"),
            expires_at=timezone.now() + timedelta(seconds=response.get("expires_in")),
        )
        return success_response(grant)
    except requests.HTTPError as e:
        logger.error(
            f"TiDB Cloud OAuth complete failed: {e.response.status_code} {e.response.text}"
        )
        return error_response("TiDB Cloud OAuth complete failed")


LAB_REGION_MAPPING = {
    ("singapore", "aws"): "regions/aws-ap-southeast-1",
    ("tokyo", "aws"): "regions/aws-ap-northeast-1",
    ("oregon", "aws"): "regions/aws-us-west-2",
    ("frankfurt", "aws"): "regions/aws-eu-central-1",
    ("singapore", "alicloud"): "regions/alicloud-ap-southeast-1",
}


def get_serverless_cluster_region(region: str, cloud_provider: str) -> str:
    return LAB_REGION_MAPPING.get((region, cloud_provider), "regions/aws-us-west-2")


def handle_http_errors(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except requests.HTTPError as err:
            match err.response.status_code:
                case 400:
                    logger.info(
                        f"TiDB Cloud API error: {err.response.status_code} {err.response.text}"
                    )
                case 401:
                    raise TiDBCloudUnauthorizedException(
                        "Unauthorized, this grant is expired or invalid, "
                        "please create a new grant and try again."
                    )
                case 429:
                    raise TiDBCloudRateLimitException(
                        "TiDB Cloud API rate limit exceeded, please try again later."
                    )
                case _:
                    pass
            raise
    return wrapper


class TiDBCloudOAuthClient:
    def __init__(
        self,
        grant: LabsOAuthTiDBCloudGrant,
    ):
        self.grant = grant
        self.request_session = requests.Session()
        retries = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        self.request_session.mount("https://", HTTPAdapter(max_retries=retries))

    def whoami(self) -> dict:
        url = "https://oauth.tidbcloud.com/v1/userinfo"
        response = self.do_request("GET", url).json()
        return response

    def create_lab_cluster(
        self,
        cluster_id: Optional[str] = None,
        region: str = "oregon",
        cloud_provider: str = "alicloud",
    ) -> UserServerlessCluster:
        """
        Create a lab cluster for the user.
        If cluster_id is provided, the cluster will be created from the existing cluster.
        If cluster_id is not provided, the cluster will be created from scratch.

        Args:
            region: The region to create the cluster in.
            cloud_provider: The cloud provider to create the cluster in.
            cluster_id: The cluster id to create the cluster from.

        Returns:
            The user serverless cluster object.
        """
        if UserServerlessCluster.objects.filter(oauth_grant=self.grant).exists():
            raise TiDBCloudServerlessException("User serverless cluster already exists")

        if cluster_id:
            try:
                cluster = self.get_serverless_cluster(cluster_id)
            except requests.HTTPError as e:
                if e.response.status_code == 404:
                    raise TiDBCloudServerlessException("Cluster not found")
                else:
                    raise TiDBCloudAPIException(
                        "TiDB Cloud OAuth client get serverless cluster failed: "
                        f"{e.response.status_code} {e.response.text}"
                    ) from e
        else:
            try:
                cluster = self.create_cluster(
                    get_serverless_cluster_region(region, cloud_provider)
                )
            except requests.HTTPError as e:
                raise TiDBCloudAPIException(
                    "TiDB Cloud OAuth client create serverless cluster failed: "
                    f"{e.response.status_code} {e.response.text}"
                ) from e

        public_disabled = cluster["endpoints"]["public"]["disabled"]
        if public_disabled:
            raise TiDBCloudServerlessException(
                "This cluster does not allow public endpoint access, please select a different one."
            )
        organization_id = cluster["labels"].get("tidb.cloud/organization")
        project_id = cluster["labels"].get("tidb.cloud/project")
        cluster_state = cluster["state"]
        cluster_id = cluster["clusterId"]
        cluster_name = cluster["displayName"]
        user_prefix = cluster["userPrefix"]
        host = cluster["endpoints"]["public"]["host"]
        port = cluster["endpoints"]["public"]["port"]
        if cluster_state == "ACTIVE":
            username, password = self.create_cluster_sql_user(cluster_id, user_prefix)
            return UserServerlessCluster.objects.create(
                user=self.grant.user,
                oauth_grant=self.grant,
                tidbcloud_organization_id=organization_id,
                tidbcloud_project_id=project_id,
                tidbcloud_cluster_id=cluster_id,
                tidbcloud_cluster_name=cluster_name,
                host=host,
                port=port,
                username=username,
                password=password,
                status=UserServerlessCluster.Status.ACTIVE,
            )
        else:
            return UserServerlessCluster.objects.create(
                user=self.grant.user,
                oauth_grant=self.grant,
                tidbcloud_organization_id=organization_id,
                tidbcloud_project_id=project_id,
                tidbcloud_cluster_id=cluster_id,
                tidbcloud_cluster_name=cluster_name,
                status=UserServerlessCluster.Status.CREATING,
            )

    def check_and_update_serverless_cluster(
        self, user_serverless_cluster: UserServerlessCluster
    ) -> UserServerlessCluster:
        """
        Check if the serverless cluster is active.
        If the cluster is ready, create a sql user and update the user serverless cluster object.

        Args:
            user_serverless_cluster: The user serverless cluster object.

        Returns:
            The updated user serverless cluster object.
        """
        if user_serverless_cluster.status == UserServerlessCluster.Status.CREATING:
            try:
                cluster = self.get_serverless_cluster(
                    user_serverless_cluster.tidbcloud_cluster_id
                )
            except requests.HTTPError as e:
                if e.response.status_code == 404:
                    raise TiDBCloudServerlessException("Cluster not found")
                else:
                    raise TiDBCloudAPIException(
                        "TiDB Cloud OAuth client get serverless cluster failed: "
                        f"{e.response.status_code} {e.response.text}"
                    ) from e
            if cluster["state"] == "ACTIVE":
                user_prefix = cluster["userPrefix"]
                host = cluster["endpoints"]["public"]["host"]
                port = cluster["endpoints"]["public"]["port"]
                username, password = self.create_cluster_sql_user(
                    user_serverless_cluster.tidbcloud_cluster_id,
                    user_prefix,
                )
                user_serverless_cluster.status = UserServerlessCluster.Status.ACTIVE
                user_serverless_cluster.username = username
                user_serverless_cluster.password = password
                user_serverless_cluster.host = host
                user_serverless_cluster.port = port
                user_serverless_cluster.save()
        return user_serverless_cluster

    def cleanup_lab_cluster(
        self, user_serverless_cluster: UserServerlessCluster
    ) -> None:
        """
        Cleanup the lab cluster.
        - Delete the sql user
        - Revoke the oauth token
        """
        try:
            self.get_serverless_cluster(user_serverless_cluster.tidbcloud_cluster_id)
        except TiDBCloudUnauthorizedException:
            self.refresh_oauth_access_token()

        self.delete_cluster_sql_user(
            user_serverless_cluster.tidbcloud_cluster_id,
            user_serverless_cluster.username,
        )
        user_serverless_cluster.sql_user_deleted_at = timezone.now()
        user_serverless_cluster.save()
        self.revoke_oauth_token()

    def list_serverless_clusters(
        self,
        recommend_region: Optional[str] = None,
        recommend_cloud_provider: Optional[str] = None,
    ) -> list[dict]:
        if recommend_region:
            recommend_region = recommend_region.lower()
        if recommend_cloud_provider:
            recommend_cloud_provider = recommend_cloud_provider.lower()
        url = "https://serverless.tidbapi.com/v1beta1/clusters"
        response = self.do_request("GET", url).json()
        clusters = response["clusters"]
        data = []
        for cluster in clusters:
            cluster_region = cluster["region"]["displayName"]
            cluster_cloud_provider = cluster["region"]["cloudProvider"]
            recommend = self._is_cluster_recommended(
                cluster_region=cluster_region,
                cluster_cloud_provider=cluster_cloud_provider,
                recommend_lab_region=recommend_region,
                recommend_cloud_provider=recommend_cloud_provider,
            )

            data.append(
                {
                    "cluster_id": cluster["clusterId"],
                    "display_name": cluster["displayName"],
                    "region": cluster_region,
                    "cloud_provider": cluster_cloud_provider,
                    "recommend": recommend,
                }
            )
        return data

    def _is_cluster_recommended(
        self,
        cluster_region: str,
        cluster_cloud_provider: str,
        recommend_lab_region: Optional[str] = None,
        recommend_cloud_provider: Optional[str] = None,
    ) -> bool:
        if not recommend_lab_region and not recommend_cloud_provider:
            return False

        region_match = True
        provider_match = True

        if recommend_lab_region:
            region_match = recommend_lab_region in cluster_region.lower()
        if recommend_cloud_provider:
            provider_match = recommend_cloud_provider == cluster_cloud_provider.lower()

        return region_match and provider_match

    def get_serverless_cluster(self, cluster_id: str) -> dict:
        url = f"https://serverless.tidbapi.com/v1beta1/clusters/{cluster_id}"
        response = self.do_request("GET", url).json()
        return response

    def create_cluster(self, region: str) -> dict:
        url = "https://serverless.tidbapi.com/v1beta1/clusters"
        response = self.do_request(
            "POST",
            url,
            json={
                "displayName": f"lab-demo-{self.grant.id}",
                "region": {
                    "name": region,
                },
            },
        ).json()
        return response

    def create_cluster_sql_user(
        self, cluster_id: str, user_prefix: str
    ) -> tuple[str, str]:
        url = f"https://iam.tidbapi.com/v1beta1/clusters/{cluster_id}/sqlUsers"
        username = f"{user_prefix}.lab_{self.grant.id}"
        password = get_random_string(length=16)
        self.do_request(
            "POST",
            url,
            json={
                "authMethod": "mysql_native_password",
                "builtinRole": "role_admin",
                "autoPrefix": False,
                "userName": username,
                "password": password,
            },
        )
        return username, password

    def delete_cluster_sql_user(self, cluster_id: str, username: str) -> None:
        url = (
            f"https://iam.tidbapi.com/v1beta1/clusters/{cluster_id}/sqlUsers/{username}"
        )
        resp = self.do_request("DELETE", url).json()
        if resp.get("code") == 200:
            logger.info(f"TiDB Cloud SQL user {username} deleted successfully")
        else:
            logger.error(f"TiDB Cloud SQL user {username} deletion failed: {resp}")

    def refresh_oauth_access_token(self) -> None:
        if self.grant.revoked_at:
            return
        tidbcloud_oauth = TiDBCloudOAuth2()
        resp = tidbcloud_oauth.refresh_token(token=self.grant.refresh_token)
        self.grant.access_token = resp["access_token"]
        self.grant.refresh_token = resp["refresh_token"]
        self.grant.expires_at = timezone.now() + timedelta(seconds=resp["expires_in"])
        self.grant.save()

    def revoke_oauth_token(self) -> None:
        if self.grant.revoked_at:
            return
        url = "https://oauth.tidbcloud.com/v1/revoke"
        client_id, client_secret = TiDBCloudOAuth2().get_key_and_secret()
        self.do_request(
            "POST",
            url,
            json={
                "token": self.grant.refresh_token,
                "token_type_hint": "refresh_token",
                "client_id": client_id,
                "client_secret": client_secret,
            },
        )
        self.grant.revoked_at = timezone.now()
        self.grant.save()

    @handle_http_errors
    def do_request(self, method: str, url: str, *args, **kwargs):
        headers = {
            "Authorization": f"Bearer {self.grant.access_token}",
            "Content-Type": "application/json",
        }
        response = self.request_session.request(
            method, url, headers=headers, *args, **kwargs
        )
        response.raise_for_status()
        return response