import logging
import random
from collections import defaultdict

import boto3

from labs.lab_manager_v2.tasks.base import BaseLabTask
from labs.lab_manager_v2.lab_spec import (
    LabSpec,
    CloudStackRequirement,
)
from labs.lab_manager_v2.exceptions import (
    LabTaskRecoverableError,
    LabTaskUnrecoverableError,
)
from labs.lab_manager_v2.schemas import (
    LabResource,
    EC2Instance,
)
from labs.lab_manager_v2.consts import get_region_metadata
from labs.models import (
    LabInstance as DBLabInstance,
)

logger = logging.getLogger(__name__)


class AWSCloudFormationStackCreateTask(BaseLabTask):
    """
    Create AWS CloudFormation stack for the lab.
    """

    def run(self):
        db_lab_instance = self.db_lab_instance
        lab_spec = self.lab_spec
        template_body = None
        for req in lab_spec.requirements:
            if isinstance(req, CloudStackRequirement):
                template_body = req.read_aws_cfn_template()
                break
        if not template_body:
            raise LabTaskUnrecoverableError(
                f"No cloud stack requirement found for lab {db_lab_instance.unique_tag}"
            )
        random_az_idx = random.randint(0, 2)
        lab_region = get_region_metadata(db_lab_instance.region_name)
        cfn = boto3.client("cloudformation", lab_region["region-code"])
        resp = cfn.create_stack(
            StackName=db_lab_instance.unique_tag,
            TemplateBody=template_body,
            Parameters=[
                {
                    "ParameterKey": "OwnerName",
                    "ParameterValue": db_lab_instance.unique_tag,
                },
                {
                    "ParameterKey": "OwnerEmail",
                    "ParameterValue": db_lab_instance.user.email,
                },
                {
                    "ParameterKey": "AZIdx",
                    "ParameterValue": str(random_az_idx),
                },
            ],
            Capabilities=["CAPABILITY_IAM"],
            OnFailure="DO_NOTHING",
        )
        logger.info(
            f"Created AWS CloudFormation stack for lab {db_lab_instance.unique_tag}, resp: {resp}"
        )
        if resp.get("ResponseMetadata", {}).get("HTTPStatusCode") != 200:
            raise LabTaskUnrecoverableError(
                f"Failed to create AWS CloudFormation stack for lab {db_lab_instance.unique_tag}"
            )


class AWSCloudFormationStackReadyCheckTask(BaseLabTask):
    """
    Check if the AWS CloudFormation stack is ready.
    """

    max_retries: int = 90
    retry_interval_seconds: int = 20

    def run(self):
        db_lab_instance = self.db_lab_instance
        lab_region = get_region_metadata(db_lab_instance.region_name)
        cfn = boto3.client("cloudformation", lab_region["region-code"])
        resp = cfn.describe_stacks(StackName=db_lab_instance.unique_tag)
        if resp.get("ResponseMetadata", {}).get("HTTPStatusCode") != 200:
            raise LabTaskUnrecoverableError(
                f"Failed to check if AWS CloudFormation stack for lab {db_lab_instance.unique_tag} is ready"
            )
        stacks = resp.get("Stacks", [])
        if len(stacks) == 0:
            raise LabTaskRecoverableError(
                f"AWS CloudFormation stack for lab {db_lab_instance.unique_tag} is not found"
            )
        stack_status = stacks[0].get("StackStatus")
        if stack_status == "CREATE_FAILED":
            raise LabTaskUnrecoverableError(
                f"AWS CloudFormation stack for lab {db_lab_instance.unique_tag} creation failed, "
                "please check the stack detail in AWS console"
            )
        elif stack_status != "CREATE_COMPLETE":
            raise LabTaskRecoverableError(
                f"AWS CloudFormation stack for lab {db_lab_instance.unique_tag} is not ready"
            )


class TagAWSEC2InstanceTask(BaseLabTask):
    """
    Tag AWS EC2 instances with the role.
    """

    max_retries: int = 20
    retry_interval_seconds: int = 20

    def run(self):
        db_lab_instance = self.db_lab_instance
        lab_spec = self.lab_spec
        lab_region = get_region_metadata(db_lab_instance.region_name)
        ec2_client = boto3.client("ec2", lab_region["region-code"])
        self.assign_roles_to_ec2_instances(db_lab_instance, lab_spec, ec2_client)

    def assign_roles_to_ec2_instances(
        self,
        db_lab_instance: DBLabInstance,
        lab_spec: LabSpec,
        ec2_client: boto3.client,
    ) -> None:
        """
        Check if the actual EC2 instance count matches the lab_spec configuration.
        """
        response = ec2_client.describe_instances(
            Filters=[
                {
                    "Name": "tag:trainer",
                    "Values": [db_lab_instance.unique_tag],
                },
            ]
        )
        if response.get("ResponseMetadata", {}).get("HTTPStatusCode") != 200:
            raise LabTaskRecoverableError(
                f"Failed to list AWS EC2 instances for lab {db_lab_instance.unique_tag}"
            )

        ec2_role_map = defaultdict(list)
        for reservation in response.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                instance_role = None
                for tag in instance.get("Tags", []):
                    if tag.get("Key") == "Name":
                        instance_role = tag.get("Value")
                        break

                if instance_role:
                    ec2_role_map[instance_role].append(
                        {
                            "instance_id": instance["InstanceId"],
                            "public_ip": instance["PublicIpAddress"],
                            "private_ip": instance["PrivateIpAddress"],
                        }
                    )

        ec2_instances = []
        for node in lab_spec.topology:
            expected_count = node.node_count
            actual_count = len(ec2_role_map.get(node.role, []))

            if actual_count != expected_count:
                raise LabTaskUnrecoverableError(
                    f"EC2 instance count mismatch: {node.role}, expected: {expected_count}, actual: {actual_count}"
                )

            for index, instance in enumerate(ec2_role_map[node.role]):
                node_role = f"{node.role}{index + 1}"
                ec2_client.create_tags(
                    Resources=[
                        instance["instance_id"],
                    ],
                    Tags=[
                        {
                            "Key": "role",
                            "Value": node_role,
                        },
                    ],
                )
                ec2_instances.append(
                    EC2Instance(
                        role=node_role,
                        public_ip=instance["public_ip"],
                        private_ip=instance["private_ip"],
                        ssh_user=node.ssh_user,
                        web_ssh=node.web_ssh,
                        https_reverse_proxy=node.https_reverse_proxy,
                    )
                )
        # 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.ec2_instances = ec2_instances
            locked_instance.resources = lab_resource.to_db_lab_instance_resource()
            locked_instance.save(update_fields=["resources"])