import logging
import random
from collections import defaultdict

from alibabacloud_tea_openapi.models import Config as AlicloudConfig
from alibabacloud_ros20190910.client import Client as AlicloudROSClient
from alibabacloud_ros20190910.models import (
    CreateStackRequest,
    CreateStackRequestParameters,
    CreateStackRequestTags,
    ListStacksRequest,
    ListStacksRequestTag,
)
from alibabacloud_ecs20140526.client import Client as AlicloudECSClient
from alibabacloud_ecs20140526.models import (
    DescribeInstancesRequest,
    DescribeInstancesRequestTag,
)
from django.conf import settings

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

logger = logging.getLogger(__name__)


class AlicloudROSStackCreateTask(BaseLabTask):
    """
    Create Alicloud ROS 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_alicloud_ros_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, db_lab_instance.cloud_provider)
        region_code = lab_region["region-code"]
        client = AlicloudROSClient(
            config=AlicloudConfig(
                access_key_id=settings.ALICLOUD_ROS_ACCESS_KEY_ID,
                access_key_secret=settings.ALICLOUD_ROS_ACCESS_KEY_SECRET,
                region_id=region_code,
            )
        )
        response = client.create_stack(
            request=CreateStackRequest(
                region_id=region_code,
                stack_name=db_lab_instance.unique_tag,
                template_body=template_body,
                parameters=[
                    CreateStackRequestParameters(
                        parameter_key="OwnerName",
                        parameter_value=db_lab_instance.unique_tag,
                    ),
                    CreateStackRequestParameters(
                        parameter_key="OwnerEmail",
                        parameter_value=db_lab_instance.user.email,
                    ),
                    CreateStackRequestParameters(
                        parameter_key="AZIdx",
                        parameter_value=str(random_az_idx),
                    ),
                ],
                tags=[
                    CreateStackRequestTags(
                        key="trainer",
                        value=db_lab_instance.unique_tag,
                    ),
                ],
            )
        )
        if response.status_code != 200:
            raise LabTaskUnrecoverableError(
                f"Failed to create Alicloud ROS stack for lab {db_lab_instance.unique_tag}, "
                f"response: {response}"
            )


class AlicloudROSStackReadyCheckTask(BaseLabTask):
    """
    Check if the Alicloud ROS stack is ready.
    """

    max_retries: int = 50
    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)
        region_code = lab_region["region-code"]
        client = AlicloudROSClient(
            config=AlicloudConfig(
                access_key_id=settings.ALICLOUD_ROS_ACCESS_KEY_ID,
                access_key_secret=settings.ALICLOUD_ROS_ACCESS_KEY_SECRET,
                region_id=region_code,
            )
        )
        response = client.list_stacks(
            request=ListStacksRequest(
                region_id=region_code,
                tag=[
                    ListStacksRequestTag(
                        key="trainer",
                        value=db_lab_instance.unique_tag,
                    )
                ],
            )
        )
        if response.status_code != 200:
            raise LabTaskRecoverableError(
                f"Failed to check if Alicloud ROS stack for lab {db_lab_instance.unique_tag} is ready, "
                f"response status code: {response.status_code}"
            )
        stacks = response.body.stacks
        if len(stacks) != 1:
            raise LabTaskRecoverableError(
                f"Failed to check if Alicloud ROS stack for lab {db_lab_instance.unique_tag} is ready, "
                f"expected 1 stack, but got {len(stacks)}"
            )
        stack = stacks[0]
        if stack.status == "CREATE_FAILED":
            raise LabTaskUnrecoverableError(
                f"Alicloud ROS stack for lab {db_lab_instance.unique_tag} creation failed, "
                "check the stack detail in Alicloud console"
            )
        elif stack.status != "CREATE_COMPLETE":
            raise LabTaskRecoverableError(
                f"Alicloud ROS stack for lab {db_lab_instance.unique_tag} is not ready, "
                f"stack status: {stack.status}"
            )


class AlicloudTagECSInstanceTask(BaseLabTask):
    """
    Tag Alicloud ECS instances with the role.
    """

    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)
        region_code = lab_region["region-code"]

        client = AlicloudECSClient(
            config=AlicloudConfig(
                access_key_id=settings.ALICLOUD_ROS_ACCESS_KEY_ID,
                access_key_secret=settings.ALICLOUD_ROS_ACCESS_KEY_SECRET,
                region_id=region_code,
            )
        )
        response = client.describe_instances(
            request=DescribeInstancesRequest(
                region_id=region_code,
                tag=[
                    DescribeInstancesRequestTag(
                        key="trainer",
                        value=db_lab_instance.unique_tag,
                    )
                ],
                page_size=100,
            )
        )
        instances = response.body.instances.instance
        ec2_role_map = defaultdict(list)
        for instance in instances:
            for tag in instance.tags.tag:
                if tag.tag_key == "Name":
                    ec2_role_map[tag.tag_value].append(
                        {
                            "public_ip": instance.public_ip_address.ip_address[0],
                            "private_ip": instance.vpc_attributes.private_ip_address.ip_address[
                                0
                            ],
                        }
                    )

        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"Alicloud ECS 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_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
        from labs.models import LabInstance as DBLabInstance

        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"])