import re
import logging
from datetime import timedelta
from typing import Optional
from urllib.parse import urlencode, urlunparse
from dataclasses import dataclass

import jwt
from django.db import transaction, models
from django.utils.translation import gettext as _
from django.utils import timezone
from django.urls import reverse
from django.utils.functional import cached_property
from django.conf import settings

from labs.models import (
    User as DBUser,
    LabInstance as DBLabInstance,
    LabJob as DBLabJob,
    LabTask as DBLabTask,
    UserServerlessCluster as DBUserServerlessCluster,
    UserLabRedemption as DBUserLabRedemption,
    UserLabWorkshop as DBUserLabWorkshop,
    LabManagementSSHKey,
    LabSSHServer,
)
from labs.billing.lab import redeem_lab_with_credits
from labs.billing.exceptions import InsufficientCreditsError
from labs.lab_manager_v2.lab_spec import (
    load_lab_spec,
    TiDBEnvironmentType,
    LabSpec,
    LabNotFoundError,
)
from labs.lab_manager_v2.schemas import LabResource
from labs.lab_manager_v2.lab_job import CeleryLabCreateJob, CeleryLabTerminateJob
from labs.lab_manager_v2.misc import record_lab_start
from labs.lab_manager_v2.exceptions import (
    LabRedeemFailedError,
    LabAccessForbiddenError,
    LabExitForbiddenError,
)
from labs.lab_manager_v2.consts import get_region_metadata
from labs.user_events import UserEvent, log_user_event

logger = logging.getLogger(__name__)


def _sanitize_string(s: str) -> str:
    return re.sub(r"[^a-zA-Z0-9]", " ", s).replace(" ", "-")


LAB_UNIQUE_TAG_MAX_LENGTH = 32

assert LAB_UNIQUE_TAG_MAX_LENGTH >= 20, (
    "lab_instance_id is a bigint, need extra space for prefixes"
)


def get_lab_unique_tag(lab_instance_id: int, email: str, lab_code: str) -> str:
    """
    The `lab_unique_tag` is used for naming lab resources in AWS. It must:
    1. Exclude special characters.
    2. Be no longer than LAB_UNIQUE_TAG_MAX_LENGTH characters.
    3. Be unique for each lab instance.
    4. Be human readable.

    `lab_instance_id` ensures uniqueness; other parts enhance tag readability and simplify debugging.
    """
    email_part = _sanitize_string(email)[:25]
    lab_code_part = _sanitize_string(lab_code)[:10]
    tag = f"{email_part}-{lab_code_part}-{lab_instance_id}"

    if len(tag) > LAB_UNIQUE_TAG_MAX_LENGTH:
        tag = tag[-LAB_UNIQUE_TAG_MAX_LENGTH:]

    if not tag[0].isalpha():
        tag = "x" + tag[1:]
    return tag


class LabInstance:
    def __init__(self, db_lab_instance: DBLabInstance):
        self.db_lab_instance = db_lab_instance
        self.lab_spec = load_lab_spec(db_lab_instance.lab_code)
        self.create_job = CeleryLabCreateJob(db_lab_instance.create_job)
        self.terminate_job = CeleryLabTerminateJob(db_lab_instance.terminate_job)

    @classmethod
    def start_new_lab(
        cls,
        *,
        user: DBUser,
        lab_spec: LabSpec,
        cloud_provider: str,
        region_name: str,
        tidb_version: str,
        tidb_environment: TiDBEnvironmentType,
        user_serverless_cluster: Optional[DBUserServerlessCluster],
    ) -> "LabInstance":
        lab_code = lab_spec.lab_code
        can_start_response = check_user_can_start_lab(user, lab_spec)
        if not can_start_response.can_start:
            raise LabAccessForbiddenError(can_start_response.suggestion)

        access_type = can_start_response.access_type

        try:
            with transaction.atomic():
                if access_type == DBLabInstance.LabAccessType.CREDITS:
                    user_lab_redemption = (
                        DBUserLabRedemption.objects.select_for_update().get(
                            id=can_start_response.user_redemption.id,
                            user=user,
                            lab_code=lab_code,
                            lab_instance__isnull=True,
                        )
                    )
                    access_detail = {
                        "redemption_id": user_lab_redemption.id,
                    }
                elif access_type == DBLabInstance.LabAccessType.WORKSHOP:
                    access_detail = {
                        "workshop_name": can_start_response.user_workshop.workshop_name,
                    }
                else:
                    access_detail = {}

                # add extra 5 minutes to cover lab startup time
                scheduled_terminate_at = timezone.now() + timedelta(
                    minutes=(lab_spec.estimated_duration_minutes + 5)
                )

                db_lab_instance = DBLabInstance.objects.create(
                    user=user,
                    lab_code=lab_code,
                    cloud_provider=cloud_provider,
                    region_name=region_name,
                    tidb_version=tidb_version,
                    tidb_environment=tidb_environment.value,
                    access_type=access_type,
                    access_detail=access_detail,
                    user_serverless_cluster=user_serverless_cluster,
                    scheduled_terminate_at=scheduled_terminate_at,
                )

                if access_type == DBLabInstance.LabAccessType.CREDITS:
                    user_lab_redemption.lab_instance = db_lab_instance
                    user_lab_redemption.save()
        except Exception:
            import traceback

            logger.error(traceback.format_exc())
            raise LabAccessForbiddenError(
                _(
                    "Failed to create lab instance, please refresh the page and try again."
                )
            )

        user.last_lab_code = lab_code
        user.last_region_name = region_name
        user.last_tidb_version = tidb_version
        user.save(
            update_fields=["last_lab_code", "last_region_name", "last_tidb_version"]
        )

        with transaction.atomic():
            create_job = CeleryLabCreateJob.create(db_lab_instance, lab_spec)
            terminate_job = CeleryLabTerminateJob.create(db_lab_instance, lab_spec)
            db_lab_instance.create_job = create_job.db_lab_job
            db_lab_instance.terminate_job = terminate_job.db_lab_job
            db_lab_instance.unique_tag = get_lab_unique_tag(
                db_lab_instance.id,
                db_lab_instance.user.email,
                db_lab_instance.lab_code,
            )
            db_lab_instance.save(
                update_fields=["create_job", "terminate_job", "unique_tag"]
            )

        record_lab_start(user.email, lab_code)
        log_user_event(
            user=user,
            event_name=UserEvent.EventName.START_LAB,
            lab_instance=db_lab_instance,
        )

        create_job.schedule_next_task()
        return cls(db_lab_instance)

    def exit_lab(self) -> None:
        """
        Exit the lab, terminate the lab instance if needed.
        User MUST call this method to exit the lab before starting a new one.
        """
        logger.info(f"User is exiting lab {self.db_lab_instance.id}")

        if self.db_lab_instance.exited_at:
            logger.info(f"Lab {self.db_lab_instance.id} already exited, skipping exit")
            return

        user_lab_status = self.db_lab_instance.get_user_status()
        if user_lab_status not in [
            DBLabInstance.LabInstanceStatus.RUNNING,
            DBLabInstance.LabInstanceStatus.CREATING_FAILED,
            DBLabInstance.LabInstanceStatus.TIMEOUT,
        ]:
            logger.info(
                f"Lab {self.db_lab_instance.id} cannot be exited in status {user_lab_status}"
            )
            raise LabExitForbiddenError(
                _("Cannot exit lab in status %s") % user_lab_status
            )

        self.terminate_lab()
        self.db_lab_instance.exited_at = timezone.now()
        self.db_lab_instance.save(update_fields=["exited_at"])

        log_user_event(
            user=self.db_lab_instance.user,
            event_name=UserEvent.EventName.TERMINATE_LAB,
            lab_instance=self.db_lab_instance,
        )

    def terminate_lab(self) -> None:
        """
        Terminate the lab instance, will clean up all the resources.
        """
        logger.info(f"Starting to terminate lab {self.db_lab_instance.id}")

        if self.db_lab_instance.status != DBLabInstance.LabInstanceStatus.RUNNING:
            logger.info(
                f"Lab {self.db_lab_instance.id} is not running, skipping termination"
            )
            return

        self.db_lab_instance.status = DBLabInstance.LabInstanceStatus.TERMINATING
        self.db_lab_instance.start_terminating_at = timezone.now()
        self.db_lab_instance.save(update_fields=["status", "start_terminating_at"])
        self.terminate_job.schedule_next_task()

    def get_job_manager(
        self, db_lab_job: DBLabJob
    ) -> CeleryLabCreateJob | CeleryLabTerminateJob:
        if db_lab_job.type == DBLabJob.TypeChoices.CREATE:
            return CeleryLabCreateJob(db_lab_job)
        elif db_lab_job.type == DBLabJob.TypeChoices.TERMINATE:
            return CeleryLabTerminateJob(db_lab_job)
        else:
            raise ValueError(f"Unsupported lab job type: {db_lab_job.type}")

    def get_environment_variables(self) -> dict:
        lab_code = self.db_lab_instance.lab_code
        region_name = self.db_lab_instance.region_name
        cloud_provider = self.db_lab_instance.cloud_provider
        region_metadata = get_region_metadata(region_name, cloud_provider)
        https_reverse_proxy_url = reverse(
            "labs-instance-https-reverse-proxy",
            kwargs={"lab_instance_id": self.db_lab_instance.id},
        )
        vars = {
            "LAB_CODE": lab_code,
            "LAB_INSTANCE_ID": self.db_lab_instance.id,
            "CLOUD_PROVIDER": cloud_provider,
            "REGION_NAME": region_name,
            "REGION_CODE": region_metadata["region-code"],
            "TIDB_VERSION": self.db_lab_instance.tidb_version,
            "TIDB_ENVIRONMENT": self.db_lab_instance.tidb_environment,
            "USER_UNIQUE_TAG": self.db_lab_instance.unique_tag,
            "HTTPS_REVERSE_PROXY_URL": https_reverse_proxy_url,
        }
        lab_resource = self.lab_resource
        for instance in lab_resource.ec2_instances:
            role = instance.role.replace("-", "_").upper()
            vars[f"HOST_{role}_PUBLIC_IP"] = instance.public_ip
            vars[f"HOST_{role}_PRIVATE_IP"] = instance.private_ip
        if lab_resource.tidb_serverless:
            tidb_serverless = lab_resource.tidb_serverless
            vars.update(
                {
                    "SERVERLESS_CLUSTER_ENVIRONMENT": tidb_serverless.environment,
                    "SERVERLESS_CLUSTER_NAME": tidb_serverless.cluster_name,
                    "SERVERLESS_CLUSTER_HOST": tidb_serverless.cluster_host,
                    "SERVERLESS_CLUSTER_PORT": tidb_serverless.cluster_port,
                    "SERVERLESS_CLUSTER_USERNAME": tidb_serverless.admin_name,
                    "SERVERLESS_CLUSTER_PASSWORD": tidb_serverless.admin_password,
                    "SERVERLESS_CLUSTER_DATABASE_NAME": tidb_serverless.database_name,
                    "SERVERLESS_CLUSTER_ID": tidb_serverless.cluster_id,
                    "SERVERLESS_CLUSTER_ORGANIZATION_ID": tidb_serverless.organization_id,
                    "SERVERLESS_CLUSTER_PROJECT_ID": tidb_serverless.project_id,
                }
            )
        if lab_resource.aws_aurora_serverless:
            aws_aurora_serverless = lab_resource.aws_aurora_serverless
            vars.update(
                {
                    "AWS_AURORA_SERVERLESS_CLUSTER_IDENTIFIER": aws_aurora_serverless.cluster_identifier,
                    "AWS_AURORA_SERVERLESS_INSTANCE_IDENTIFIER": aws_aurora_serverless.instance_identifier,
                    "AWS_AURORA_SERVERLESS_PORT": aws_aurora_serverless.port,
                    "AWS_AURORA_SERVERLESS_ENDPOINT": aws_aurora_serverless.endpoint,
                    "AWS_AURORA_SERVERLESS_READER_ENDPOINT": aws_aurora_serverless.reader_endpoint,
                    "AWS_AURORA_SERVERLESS_DATABASE_NAME": aws_aurora_serverless.database_name,
                    "AWS_AURORA_SERVERLESS_MASTER_USERNAME": aws_aurora_serverless.master_username,
                    "AWS_AURORA_SERVERLESS_MASTER_PASSWORD": aws_aurora_serverless.master_password,
                }
            )
        if self.db_lab_instance.user_serverless_cluster:
            vars["TIDBCLOUD_ACCESS_TOKEN"] = (
                self.db_lab_instance.user_serverless_cluster.oauth_grant.access_token
            )
        return vars

    @cached_property
    def lab_resource(self) -> LabResource:
        return LabResource.from_db_lab_instance_resource(self.db_lab_instance.resources)

    def get_https_reverse_proxy_url(
        self, redirect_to: str, machine_role: str
    ) -> Optional[str]:
        ssh_jump_ip = None
        for ei in self.lab_resource.ec2_instances:
            if ei.role == machine_role:
                ssh_jump_ip = ei.public_ip
                break

        if not ssh_jump_ip:
            logger.error(
                f"Machine role {machine_role} not found in lab instance "
                f"{self.db_lab_instance.id} when getting https reverse proxy url"
            )
            return None

        payload = {
            "exp": int(timezone.now().timestamp() + 3600 * 2),
            "iss": "https://labs.tidb.io",
            "sub": str(self.db_lab_instance.user_id),
            "uid": str(self.db_lab_instance.user_id),
            "aud": str(self.db_lab_instance.id),
        }
        # decode the private key from unicode_escape to utf-8
        # because the private key is encoded in unicode_escape in the aws secrets manager
        jwt_private_key = settings.WEB_SSH_JWT_PRIVATE_KEY.encode("utf-8").decode(
            "unicode_escape"
        )
        token = jwt.encode(payload, jwt_private_key, algorithm="EdDSA")
        reverse_proxy_host = (
            ssh_jump_ip.replace(".", "-") + "." + self.db_lab_instance.xip_domain
        )
        query = urlencode(
            {
                "user_session": token,
                "redirect": redirect_to or "/",
            }
        )
        url = urlunparse(("https", reverse_proxy_host, "/set-token", "", query, ""))
        return url

    def get_secure_access_proxy_url(
        self, redirect_to: str, machine_role: str
    ) -> Optional[str]:
        return self.get_https_reverse_proxy_url(redirect_to, machine_role)

    def get_management_ssh_key(self) -> LabManagementSSHKey:
        return LabManagementSSHKey.objects.get(
            region_name=self.db_lab_instance.region_name,
            cloud_provider=self.db_lab_instance.cloud_provider,
            is_active=True,
        )

    def get_detail(self) -> dict:
        """
        Get the detail of the lab instance for user,
        hide some fields that user should not see and care about.
        """
        creation_progress = None
        estimated_creation_duration_seconds = None
        now = timezone.now()
        lab_remaining_seconds = (
            self.db_lab_instance.scheduled_terminate_at - now
        ).total_seconds()

        if self.db_lab_instance.status == DBLabInstance.LabInstanceStatus.CREATING:
            creation_progress = {
                "total": DBLabTask.objects.filter(
                    lab_job=self.db_lab_instance.create_job
                ).count(),
                "completed": DBLabTask.objects.filter(
                    lab_job=self.db_lab_instance.create_job,
                    status=DBLabTask.StatusChoices.SUCCEEDED,
                ).count(),
            }
            estimated_creation_duration_seconds = get_lab_ready_duration_seconds(
                self.db_lab_instance.lab_code
            )

        return {
            "id": self.db_lab_instance.id,
            "lab_code": self.db_lab_instance.lab_code,
            "cloud_provider": self.db_lab_instance.cloud_provider,
            "region_name": self.db_lab_instance.region_name,
            "tidb_version": self.db_lab_instance.tidb_version,
            "tidb_environment": self.db_lab_instance.tidb_environment,
            "status": self.db_lab_instance.get_user_status(),
            "created_at": self.db_lab_instance.created_at,
            "ready_at": self.db_lab_instance.ready_at,
            "exited_at": self.db_lab_instance.exited_at,
            "creation_progress": creation_progress,
            "estimated_creation_duration_seconds": estimated_creation_duration_seconds,
            "remaining_seconds": max(int(lab_remaining_seconds), 0),
        }

    def get_resources(self) -> list[dict]:
        """
        Build the lab resources list for frontend, including:
        - EC2 instances
        - TiDB Cloud Serverless
        - TiDB Dashboard
        """
        data = []
        if self.db_lab_instance.status != DBLabInstance.LabInstanceStatus.RUNNING:
            return data

        resource = LabResource.from_db_lab_instance_resource(
            self.db_lab_instance.resources
        )
        data.extend(resource.to_frontend())

        for ei in resource.ec2_instances:
            if ei.web_ssh:
                web_ssh_url = self._get_web_ssh_url(
                    ei.ssh_user,
                    ei.public_ip,
                )
                data.append(
                    {
                        "Role": f"{ei.role}-ssh",
                        "Type": "Web SSH",
                        "Endpoint": web_ssh_url,
                    }
                )
        return data

    def _get_web_ssh_url(self, ssh_user: str, public_ip: str) -> str:
        payload = {
            "id": self.db_lab_instance.id,
            "user_id": self.db_lab_instance.user.id,
            "lab_code": self.db_lab_instance.lab_code,
            "cloud_provider": self.db_lab_instance.cloud_provider,
            "region": self.db_lab_instance.region_name,
            "m_ip": public_ip,
            "m_port": 22,
            "m_user": ssh_user,
            "exp": int(timezone.now().timestamp() + 3600 * 2),
        }
        # decode the private key from unicode_escape to utf-8
        # because the private key is encoded in unicode_escape in the aws secrets manager
        jwt_private_key = settings.WEB_SSH_JWT_PRIVATE_KEY.encode("utf-8").decode(
            "unicode_escape"
        )
        token = jwt.encode(payload, jwt_private_key, algorithm="EdDSA")
        if self._web_ssh_server:
            return f"wss://{self._web_ssh_server}/ssh-server/web?auth={token}"
        else:
            logger.error(
                f"No available web ssh server found for lab {self.db_lab_instance.id}, "
                f"region {self.db_lab_instance.region_name}"
            )

    SSH_SERVER_ACTIVE_SESSION_THRESHOLD = 200

    @cached_property
    def _web_ssh_server(self) -> Optional[str]:
        ssh_servers = LabSSHServer.objects.filter(
            status=LabSSHServer.Status.HEALTHY,
        ).values("domain", "regions", "active_session")
        if not ssh_servers:
            return None

        if len(ssh_servers) == 1:
            return ssh_servers[0]["domain"]

        lab_region = self.db_lab_instance.region_name
        servers = sorted(ssh_servers, key=lambda s: s["active_session"])
        low_active_session_servers = [
            s
            for s in servers
            if s["active_session"] < self.SSH_SERVER_ACTIVE_SESSION_THRESHOLD
        ]
        candidates = low_active_session_servers or servers
        region_candidates = [
            s for s in candidates if lab_region in (s.get("regions") or [])
        ]
        if region_candidates:
            return region_candidates[0]["domain"]

        return candidates[0]["domain"]


def redeem_lab(user: DBUser, lab_code: str, lab_spec: LabSpec) -> None:
    """
    Redeem the lab for the user.
    """
    if not lab_spec.credit_redeemable:
        raise LabRedeemFailedError(_("Lab is not redeemable with credits"))

    if user.is_pingcap_employee:
        raise LabRedeemFailedError(_("You are a PingCAP employee, you don't need to redeem this lab with credits"))

    if lab_spec.employee_only:
        raise LabRedeemFailedError(_("This lab is for PingCAP employees only"))

    if DBUserLabRedemption.objects.filter(
        user=user,
        lab_code=lab_code,
        lab_instance__isnull=True,
    ).exists():
        raise LabRedeemFailedError(
            _(
                "You have already redeemed this lab, please refresh the page and try again."
            )
        )

    try:
        redeem_lab_with_credits(
            user=user, lab_code=lab_code, credits_to_deduct=lab_spec.credits
        )
    except InsufficientCreditsError:
        raise LabRedeemFailedError(
            _(
                "You don't have enough credits to redeem this lab. "
                "Please buy or acquire credits first."
            )
        )
    log_user_event(
        user=user,
        event_name=UserEvent.EventName.REDEEM_LAB,
        lab_code=lab_code,
    )


def get_lab_ready_duration_seconds(lab_code: str) -> int:
    # Calculate the average ready duration of the last 3 lab instances
    result = (
        DBLabInstance.objects.filter(
            lab_code=lab_code,
            created_at__isnull=False,
            ready_at__isnull=False,
        )
        .annotate(
            duration_seconds=models.ExpressionWrapper(
                models.F("ready_at") - models.F("created_at"),
                output_field=models.DurationField(),
            )
        )
        .order_by("created_at")[:3]
        .aggregate(avg_duration=models.Avg("duration_seconds"))
    )

    if not result["avg_duration"]:
        # If there are no historical data, use the default value 360 seconds
        return 360

    return int(result["avg_duration"].total_seconds())


def get_active_lab_instance(
    user: DBUser,
    lab_code: Optional[str] = None,
) -> Optional[LabInstance]:
    """
    Get the active lab instance for the user.
    """
    qs = DBLabInstance.objects.select_related("user").filter(
        user=user,
        exited_at__isnull=True,
    )
    if lab_code:
        qs = qs.filter(lab_code=lab_code)
    active_lab_instance = qs.order_by("-created_at").first()
    if not active_lab_instance:
        return None

    try:
        return LabInstance(active_lab_instance)
    except LabNotFoundError:
        logger.error(f"Lab {active_lab_instance.lab_code} not found, its spec may be deleted")
        return None


@dataclass
class CanStartLabResponse:
    can_start: bool
    suggestion: str
    access_type: DBLabInstance.LabAccessType
    user_workshop: Optional[DBUserLabWorkshop] = None
    user_redemption: Optional[DBUserLabRedemption] = None


def check_user_can_start_lab(user: DBUser, lab_spec: LabSpec) -> CanStartLabResponse:
    """
    Validates whether a user can start a new lab by checking various access conditions.

    Access checks are performed in the following priority order:
    1. Active lab check - Users cannot have multiple labs running simultaneously
    2. Employee check - PingCAP employees have unrestricted access
    3. Workshop access - Check for valid workshop enrollment
    4. Credits check - Verify if user has redeemed credits for the lab

    Args:
        user: The user attempting to start the lab
        lab_code: The identifier code for the lab to be started

    Returns:
        CanStartLabResponse
    """
    active_lab_instance = get_active_lab_instance(user)
    if active_lab_instance:
        return CanStartLabResponse(
            can_start=False,
            suggestion=_("You have an active lab in-progress."),
            access_type=DBLabInstance.LabAccessType.CREDITS,
        )

    if user.email.strip().lower().endswith("@pingcap.com"):
        return CanStartLabResponse(
            can_start=True,
            suggestion="",
            access_type=DBLabInstance.LabAccessType.INTERNAL,
        )
    else:
        if lab_spec.employee_only:
            return CanStartLabResponse(
                can_start=False,
                suggestion=_("This lab is for PingCAP employees only."),
                access_type=DBLabInstance.LabAccessType.INTERNAL,
            )

    now = timezone.now()
    user_lab_workshop = DBUserLabWorkshop.objects.filter(
        user=user,
        lab_codes__contains=lab_spec.lab_code,
        valid_from__lte=now,
        valid_until__gte=now,
    ).first()
    if user_lab_workshop:
        return CanStartLabResponse(
            can_start=True,
            suggestion="",
            access_type=DBLabInstance.LabAccessType.WORKSHOP,
            user_workshop=user_lab_workshop,
        )

    if not lab_spec.credit_redeemable:
        return CanStartLabResponse(
            can_start=False,
            suggestion=_("This lab is for instructor-led workshop only."),
            access_type=DBLabInstance.LabAccessType.WORKSHOP,
        )

    user_lab_redemption = (
        DBUserLabRedemption.objects.filter(
            user=user,
            lab_code=lab_spec.lab_code,
            lab_instance__isnull=True,
        )
        .order_by("created_at")
        .first()
    )
    if not user_lab_redemption:
        return CanStartLabResponse(
            can_start=False,
            suggestion=_("Redeem with credits to start this lab."),
            access_type=DBLabInstance.LabAccessType.CREDITS,
        )

    return CanStartLabResponse(
        can_start=True,
        suggestion="",
        access_type=DBLabInstance.LabAccessType.CREDITS,
        user_redemption=user_lab_redemption,
    )