import io
import logging
import shlex
from itertools import product

import paramiko
from scp import SCPClient
from django.conf import settings

from labs.lab_manager_v2.tasks.base import BaseLabTask
from labs.lab_manager_v2.lab_spec import (
    BootstrapShellScriptStep,
)
from labs.lab_manager_v2.exceptions import (
    LabTaskRecoverableError,
    LabTaskUnrecoverableError,
)
from labs.lab_manager_v2.schemas import (
    LabResource,
    EC2Instance,
)
from labs.models import (
    LabInstance as DBLabInstance,
    UserServerlessCluster as DBUserServerlessCluster,
    LabManagementSSHKey,
)
from labs.tidb_serverless.branch import checkout_serverless_branch
from labs.tidb_serverless.oauth import TiDBCloudOAuthClient
from labs.tidb_serverless.exceptions import TiDBCloudException
from labs.tidb_serverless.schemas import TiDBServerlessClusterConfig
from labs.utils.network import check_ec2_instance_port

logger = logging.getLogger(__name__)


class TiDBCloudServerlessClusterCreateTask(BaseLabTask):
    """
    Create TiDB Cloud Serverless cluster for the lab.
    If the lab has a user serverless cluster, use it.
    Otherwise, create a sandbox serverless cluster.
    """

    max_retries: int = 50
    retry_interval_seconds: int = 5

    def run(self):
        db_lab_instance = self.db_lab_instance
        if db_lab_instance.user_serverless_cluster:
            serverless_info = self.get_user_serverless_cluster(
                db_lab_instance.user_serverless_cluster,
                db_lab_instance,
            )
        else:
            serverless_info = self.create_sandbox_serverless_cluster()

        serverless_info.database_name = (
            f"lab_{db_lab_instance.lab_code}_{db_lab_instance.id}"
        )

        # Use select_for_update to prevent race conditions with other parallel tasks
        from django.db import transaction

        with transaction.atomic():
            locked_instance = DBLabInstance.objects.select_for_update().get(pk=db_lab_instance.pk)
            lab_resource = LabResource.from_db_lab_instance_resource(locked_instance.resources)
            lab_resource.tidb_serverless = serverless_info
            locked_instance.resources = lab_resource.to_db_lab_instance_resource()
            locked_instance.save(update_fields=["resources"])

    def get_user_serverless_cluster(
        self,
        user_serverless_cluster: DBUserServerlessCluster,
        db_lab_instance: DBLabInstance,
    ) -> TiDBServerlessClusterConfig:
        try:
            oauth_client = TiDBCloudOAuthClient(user_serverless_cluster.oauth_grant)
            oauth_client.check_and_update_serverless_cluster(user_serverless_cluster)
        except TiDBCloudException as e:
            raise LabTaskUnrecoverableError(
                f"Failed to check and update user serverless cluster for lab {db_lab_instance.unique_tag}, "
                f"error: {e}"
            ) from e

        if user_serverless_cluster.status == DBUserServerlessCluster.Status.ACTIVE:
            return TiDBServerlessClusterConfig(
                environment="user_hosted",
                cluster_id=user_serverless_cluster.tidbcloud_cluster_id,
                cluster_name=user_serverless_cluster.tidbcloud_cluster_name,
                cluster_host=user_serverless_cluster.host,
                cluster_port=user_serverless_cluster.port,
                admin_name=user_serverless_cluster.username,
                admin_password=user_serverless_cluster.password,
                organization_id=user_serverless_cluster.tidbcloud_organization_id,
                project_id=user_serverless_cluster.tidbcloud_project_id,
            )
        else:
            raise LabTaskRecoverableError(
                f"User serverless cluster for lab {db_lab_instance.unique_tag} is not ready"
            )

    def create_sandbox_serverless_cluster(self) -> TiDBServerlessClusterConfig:
        db_lab_instance = self.db_lab_instance
        serverless_branch_info = checkout_serverless_branch(
            db_lab_instance.unique_tag,
            db_lab_instance.region_name,
            db_lab_instance.cloud_provider,
        )
        if not serverless_branch_info:
            raise LabTaskUnrecoverableError(
                f"Failed to get available TiDB Cloud Serverless cluster for lab {db_lab_instance.unique_tag}"
            )
        return serverless_branch_info


class EC2ReadyCheckTask(BaseLabTask):
    """
    Check if the AWS EC2 instances are ready, which means the SSH port is accessible.
    """

    max_retries: int = 20
    retry_interval_seconds: int = 5

    def run(self):
        db_lab_instance = self.db_lab_instance
        ec2_instances = db_lab_instance.resources.get("ec2_instances", [])
        for instance in ec2_instances:
            if not check_ec2_instance_port(instance["Public IP"], 22):
                raise LabTaskRecoverableError(
                    f"EC2 instance {instance['Role']}[{instance['Public IP']}] is not accessible"
                )


def _load_management_ssh_private_key(management_ssh_key: LabManagementSSHKey):
    """Helper function to load SSH private key from LabManagementSSHKey object."""
    private_key_str = management_ssh_key.private_key
    with io.StringIO(private_key_str) as key_text_stream:
        if management_ssh_key.key_type == LabManagementSSHKey.KeyType.RSA:
            return paramiko.RSAKey.from_private_key(key_text_stream)
        elif management_ssh_key.key_type == LabManagementSSHKey.KeyType.ED25519:
            return paramiko.Ed25519Key.from_private_key(key_text_stream)
        else:
            raise LabTaskUnrecoverableError(
                f"Unsupported ssh key type: {management_ssh_key.key_type}"
            )


class ExecuteBootstrapScriptTask(BaseLabTask):
    def run(self):
        db_task = self.db_task
        db_lab_instance = self.db_lab_instance
        step_index = db_task.payload["index"]
        step: BootstrapShellScriptStep = self.lab_spec.bootstrap[step_index]
        script = io.BytesIO(step.get_shell_script().encode("utf-8"))

        from labs.lab_manager_v2.lab_instance import LabInstance

        lab_instance = LabInstance(db_lab_instance)
        lab_resource = lab_instance.lab_resource
        management_ssh_key = lab_instance.get_management_ssh_key()
        environment_variables = lab_instance.get_environment_variables()

        for target in step.target:
            for instance in lab_resource.ec2_instances:
                if instance.role == target:
                    if not instance.public_ip or not instance.ssh_user:
                        raise LabTaskUnrecoverableError(
                            f"lab {db_lab_instance.unique_tag} has no public IP or SSH user for instance {instance.role}"
                        )
                    logger.info(f"execute script on {instance.public_ip}")
                    self._execute_script(
                        instance.public_ip,
                        instance.ssh_user,
                        management_ssh_key,
                        script,
                        environment_variables,
                    )

    def _execute_script(
        self,
        ip: str,
        ssh_user: str,
        management_ssh_key: LabManagementSSHKey,
        script: io.BytesIO,
        envs: dict,
    ) -> None:
        remote_file_path = f"/tmp/bootstrap-script-{self.db_task.id}.sh"

        pkey = _load_management_ssh_private_key(management_ssh_key)
        private_key_str = management_ssh_key.private_key

        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(
            hostname=ip,
            username=ssh_user,
            pkey=pkey,
        )

        script.seek(0)
        scp = SCPClient(ssh.get_transport())
        scp.putfo(script, remote_file_path)

        ssh_private_key_tmp_file = (
            f"/tmp/pe-class-key-{self.db_lab_instance.region_name}.pem"
        )
        ssh_private_key_file_name = (
            f"~/.ssh/pe-class-key-{self.db_lab_instance.region_name}.pem"
        )
        # Upload the private key to the remote host. Use a fresh bytes stream.
        private_key_stream = io.BytesIO(private_key_str.encode("utf-8"))
        scp.putfo(private_key_stream, ssh_private_key_tmp_file)
        scp.close()

        stdin, stdout, stderr = ssh.exec_command(
            f"mv -f {ssh_private_key_tmp_file} {ssh_private_key_file_name} "
            f"&& chmod 400 {ssh_private_key_file_name}"
        )
        exit_code = stdout.channel.recv_exit_status()
        if exit_code != 0:
            raise LabTaskUnrecoverableError(
                f"Failed to chmod 400 private key on {ip}, exit_code: {exit_code}, "
                f"stdout: {stdout.read()}, stderr: {stderr.read()}, lab: {self.db_lab_instance.unique_tag}"
            )

        stdin, stdout, stderr = ssh.exec_command(f"chmod +x {remote_file_path}")
        exit_code = stdout.channel.recv_exit_status()
        if exit_code != 0:
            raise LabTaskUnrecoverableError(
                f"Failed to chmod +x script on {ip}, exit_code: {exit_code}, "
                f"stderr: {stderr.read()}, lab: {self.db_lab_instance.unique_tag}"
            )

        env_str = " ".join([f"{k}={shlex.quote(str(v))}" for k, v in envs.items()])
        stdin, stdout, stderr = ssh.exec_command(
            f"export {env_str} && {remote_file_path}"
        )
        exit_code = stdout.channel.recv_exit_status()
        if exit_code != 0:
            raise LabTaskUnrecoverableError(
                f"Failed to execute script on {ip}, exit_code: {exit_code}, "
                f"stderr: {stderr.read()}, lab: {self.db_lab_instance.unique_tag}"
            )

        stdin, stdout, stderr = ssh.exec_command(f"rm -f {remote_file_path}")
        exit_code = stdout.channel.recv_exit_status()
        if exit_code != 0:
            raise LabTaskUnrecoverableError(
                f"Failed to remove script on {ip}, exit_code: {exit_code}, "
                f"stderr: {stderr.read()}, lab: {self.db_lab_instance.unique_tag}"
            )
        ssh.close()


class SetupHTTPSReverseProxyTask(BaseLabTask):
    def run(self):
        from labs.lab_manager_v2.lab_instance import LabInstance

        db_lab_instance = self.db_lab_instance
        lab_instance = LabInstance(db_lab_instance)

        reverse_proxy_instances = []
        for ei in lab_instance.lab_resource.ec2_instances:
            if ei.https_reverse_proxy:
                if ei.ssh_user != "ubuntu":
                    raise LabTaskUnrecoverableError(
                        "HTTPS reverse proxy currently only supports Ubuntu, but "
                        f"instance {ei.role} has SSH user {ei.ssh_user}"
                    )
                reverse_proxy_instances.append(ei)

        if len(reverse_proxy_instances) == 0:
            return

        management_ssh_key = lab_instance.get_management_ssh_key()

        for instance in reverse_proxy_instances:
            self._setup_https_reverse_proxy(instance, management_ssh_key)

    AVAILABLE_XIP_DOMAINS = [
        "sslip.io",
        "nip.io",
        "xip.labs.tidb.io",
    ]
    AVAILABLE_TLS_CA = [
        "letsencrypt",
        "zerossl",
    ]

    def _setup_https_reverse_proxy(
        self, instance: EC2Instance, management_ssh_key: LabManagementSSHKey
    ) -> None:
        pkey = _load_management_ssh_private_key(management_ssh_key)

        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(
            hostname=instance.public_ip,
            username=instance.ssh_user,
            pkey=pkey,
        )
        jwt_sign_key = settings.WEB_SSH_JWT_SIGN_KEY
        for xip_domain, tls_ca in product(
            self.AVAILABLE_XIP_DOMAINS, self.AVAILABLE_TLS_CA
        ):
            domain = instance.public_ip.replace(".", "-") + "." + xip_domain
            stdin, stdout, stderr = ssh.exec_command(
                f'DOMAIN={domain} TLS_CA={tls_ca} LAB_INSTANCE_ID={self.db_lab_instance.id} JWT_SIGN_KEY="{jwt_sign_key}" '
                "bash <(curl -fsSL https://pingcap-edu-japan.s3.ap-northeast-1.amazonaws.com/caddy/install-caddy-v2.sh)"
            )
            exit_code = stdout.channel.recv_exit_status()
            if exit_code == 0:
                self.db_lab_instance.xip_domain = xip_domain
                self.db_lab_instance.save(update_fields=["xip_domain"])
                break
            logger.info(
                f"Failed to install Caddy for lab {self.db_lab_instance.unique_tag}, "
                f"xip domain: {xip_domain}, tls ca: {tls_ca}, exit_code: {exit_code}, stderr: {stderr.read()}, "
                "will try next xip domain and TLS CA."
            )
        else:
            raise LabTaskUnrecoverableError(
                f"Failed to install Caddy for lab {self.db_lab_instance.unique_tag}, "
                f"all xip domains and TLS CAs failed"
            )
        ssh.close()