# NOTE: task_id IS THE PIPELINE_ID INTERNALLY

import uuid
from typing import Any, Dict, List, Optional, Set

from cloudenv.server.providers.aws.adapter import AWSAdapter
from pydantic import BaseModel, Field


class PolicyStatement(BaseModel):
    Effect: str
    Action: List[str]
    Resource: str


class Policy(BaseModel):
    Version: str
    Statement: List[PolicyStatement]


class PipelineNode(BaseModel):
    # separate service and instance
    service: str
    instance: str
    policies: List[Policy]
    config: Dict[str, Any] = Field(default_factory=dict)


class Pipeline:
    CREATE_ACTIONS = {
        "s3": "CreateBucket",
        "lambda": "CreateFunction",
        "sqs": "CreateQueue",
        "sns": "CreateTopic",
        "dynamodb": "CreateTable",
        "iam": "CreateRole",
        "ec2": "RunInstances",
    }

    def __init__(self, pipeline_id: Optional[str] = None):
        self.pipeline_id = pipeline_id or uuid.uuid4().hex
        self.broken_graph: Dict[str, List[PipelineNode]] = {}
        self.ideal_graph: Dict[str, List[PipelineNode]] = {}

    # --------------------------------------------------
    # GRAPH BUILDING
    # --------------------------------------------------

    def add_from_json(
        self, json_data: Dict[str, List[Any]], kind: str = "broken"
    ):
        graph = self.get_graph(kind)

        for src_service, nodes in json_data.items():
            if src_service not in graph:
                graph[src_service] = []

            for node_data in nodes:
                # Extract service + instance safely
                service_full = getattr(node_data, "service", src_service)
                instance_name = getattr(node_data, "instance", None)

                if ":" in service_full:
                    service_name, parsed_instance = service_full.split(":", 1)
                    instance_name = instance_name or parsed_instance
                else:
                    service_name = service_full
                    instance_name = instance_name or service_name

                # -------------------------------
                # Policies
                # -------------------------------
                policies_raw = getattr(node_data, "policies", [])
                policies: List[Policy] = []

                for p in policies_raw:
                    if isinstance(p, dict):
                        policies.append(Policy(**p))
                    else:
                        policies.append(Policy(**p.dict()))

                # -------------------------------
                # Config
                # -------------------------------
                config = getattr(node_data, "config", {}) or {}

                # -------------------------------
                # Check existing node
                # -------------------------------
                existing_node = next(
                    (
                        n
                        for n in graph[src_service]
                        if n.service == service_name
                        and n.instance == instance_name
                    ),
                    None,
                )

                if existing_node:
                    # Merge policies
                    for new_policy in policies:
                        if new_policy not in existing_node.policies:
                            existing_node.policies.append(new_policy)

                    # Merge config
                    existing_node.config.update(config)

                else:
                    # Create new node ONLY if it doesn't exist
                    node = PipelineNode(
                        service=service_name,
                        instance=instance_name,
                        policies=policies,
                        config=config,
                    )
                    graph[src_service].append(node)

    # --------------------------------------------------
    # ACCESSORS
    # --------------------------------------------------

    def get_graph(self, kind: str) -> Dict[str, List[PipelineNode]]:
        return self.broken_graph if kind == "broken" else self.ideal_graph

    def get_node(
        self, service: str, instance: str, kind: str = "broken"
    ) -> Optional[PipelineNode]:
        graph = self.get_graph(kind)
        for nodes in graph.values():
            for node in nodes:
                if node.service == service and node.instance == instance:
                    return node
        return None

    def all_nodes(self) -> List[PipelineNode]:
        nodes: List[PipelineNode] = []
        for graph in [self.broken_graph, self.ideal_graph]:
            for _, dsts in graph.items():
                nodes.extend(dsts)
        return nodes

    # --------------------------------------------------
    # AWS MATERIALIZATION
    # --------------------------------------------------

    def build_payload(self, service: str, instance: str) -> Dict:
        if service == "s3":
            return {"Bucket": instance}
        if service == "sqs":
            return {"QueueName": instance}
        if service == "sns":
            return {"Name": instance}
        if service == "dynamodb":
            return {
                "TableName": instance,
                "KeySchema": [{"AttributeName": "id", "KeyType": "HASH"}],
                "AttributeDefinitions": [
                    {"AttributeName": "id", "AttributeType": "S"}
                ],
                "BillingMode": "PAY_PER_REQUEST",
            }
        if service == "lambda":
            return {
                "FunctionName": instance,
                "Runtime": "python3.12",
                "Role": "arn:aws:iam::123456789012:role/dummy",
                "Handler": "index.handler",
                "Code": {"ZipFile": b"dummy"},
            }
        if service == "iam":
            return {
                "RoleName": instance,
                "AssumeRolePolicyDocument": "{}",
            }
        if service == "ec2":
            return {"MinCount": 1, "MaxCount": 1}
        return {}

    def materialize(self, aws: AWSAdapter, episode_id: str) -> None:
        created: Set[str] = set()
        print("Pipeline materialization started")

        for node in self.all_nodes():

            node_key = f"{node.service}:{node.instance}"
            if node_key in created:
                continue

            action = self.CREATE_ACTIONS.get(node.service)
            if not action:
                continue

            payload = self.build_payload(node.service, node.instance)

            try:
                aws.execute(
                    episode_id=episode_id,
                    task_id=self.pipeline_id,
                    service=node.service,
                    action=action,
                    payload=payload,
                )
                created.add(node_key)

                # Attach policies to the instance
                if node.policies:
                    iam_enforcer = aws.iam_enforcer
                    policies_dicts = [p.model_dump() for p in node.policies]
                    iam_enforcer.set_instance_policies(
                        node.service, node.instance, policies_dicts
                    )

            except Exception as e:
                print("Materialization error:", e)
                # Ignore "already exists"
                pass

    # --------------------------------------------------
    # CONFIG CRUD
    # --------------------------------------------------

    def add_node_config(
        self,
        service: str,
        instance: str,
        config: Dict[str, Any],
        kind: str = "broken",
    ) -> bool:
        node = self.get_node(service, instance, kind)
        if not node:
            return False
        node.config.update(config)
        return True

    def update_node_config(
        self,
        service: str,
        instance: str,
        config: Dict[str, Any],
        kind: str = "broken",
    ) -> bool:
        return self.add_node_config(service, instance, config, kind)

    def delete_node_config(
        self,
        service: str,
        instance: str,
        keys: List[str],
        kind: str = "broken",
    ) -> bool:
        node = self.get_node(service, instance, kind)
        if not node:
            return False

        for key in keys:
            node.config.pop(key, None)

        return True

    def get_node_config(
        self, service: str, instance: str, kind: str = "broken"
    ) -> Optional[Dict[str, Any]]:
        node = self.get_node(service, instance, kind)
        return node.config if node else None

    def get_all_configs(
        self, kind: str = "broken"
    ) -> Dict[str, Dict[str, Any]]:
        graph = self.get_graph(kind)
        result = {}

        for src_service, nodes in graph.items():
            for node in nodes:
                key = f"{src_service}->{node.service}:{node.instance}"
                result[key] = node.config

        return result