from __future__ import annotations

from datetime import datetime
from typing import Any

from typing_extensions import NotRequired, TypedDict

from infrahub_sdk.graphql.query import Query


class NodeDiff(TypedDict):
    branch: str
    kind: str
    id: str
    action: str
    display_label: str
    elements: list[NodeDiffElement]


class NodeDiffElement(TypedDict):
    name: str
    element_type: str
    action: str
    summary: NodeDiffSummary
    peers: NotRequired[list[NodeDiffPeer]]


class NodeDiffSummary(TypedDict):
    added: int
    updated: int
    removed: int


class NodeDiffPeer(TypedDict):
    action: str
    summary: NodeDiffSummary


class DiffTreeData(TypedDict):
    num_added: int
    num_updated: int
    num_removed: int
    num_conflicts: int
    to_time: str
    from_time: str
    base_branch: str
    diff_branch: str
    name: NotRequired[str | None]
    nodes: list[NodeDiff]


def get_diff_summary_query() -> str:
    return """
        query GetDiffTree($branch_name: String!, $name: String, $from_time: DateTime, $to_time: DateTime) {
            DiffTree(branch: $branch_name, name: $name, from_time: $from_time, to_time: $to_time) {
                nodes {
                    uuid
                    kind
                    status
                    label
                    num_added
                    num_updated
                    num_removed
                    attributes {
                        name
                        status
                        num_added
                        num_updated
                        num_removed
                    }
                    relationships {
                        name
                        status
                        cardinality
                        num_added
                        num_updated
                        num_removed
                        elements {
                            status
                            num_added
                            num_updated
                            num_removed
                        }
                    }
                }
            }
        }
    """


def diff_tree_node_to_node_diff(node_dict: dict[str, Any], branch_name: str) -> NodeDiff:
    element_diffs: list[NodeDiffElement] = []
    if "attributes" in node_dict:
        for attr_dict in node_dict["attributes"]:
            attr_diff = NodeDiffElement(
                action=str(attr_dict.get("status")),
                element_type="ATTRIBUTE",
                name=str(attr_dict.get("name")),
                summary={
                    "added": int(attr_dict.get("num_added") or 0),
                    "removed": int(attr_dict.get("num_removed") or 0),
                    "updated": int(attr_dict.get("num_updated") or 0),
                },
            )
            element_diffs.append(attr_diff)
    if "relationships" in node_dict:
        for relationship_dict in node_dict["relationships"]:
            is_cardinality_one = str(relationship_dict.get("cardinality")).upper() == "ONE"
            relationship_diff = NodeDiffElement(
                action=str(relationship_dict.get("status")),
                element_type="RELATIONSHIP_ONE" if is_cardinality_one else "RELATIONSHIP_MANY",
                name=str(relationship_dict.get("name")),
                summary={
                    "added": int(relationship_dict.get("num_added") or 0),
                    "removed": int(relationship_dict.get("num_removed") or 0),
                    "updated": int(relationship_dict.get("num_updated") or 0),
                },
            )
            if not is_cardinality_one and "elements" in relationship_dict:
                peer_diffs = [
                    NodeDiffPeer(
                        action=str(element_dict.get("status")),
                        summary={
                            "added": int(element_dict.get("num_added") or 0),
                            "removed": int(element_dict.get("num_removed") or 0),
                            "updated": int(element_dict.get("num_updated") or 0),
                        },
                    )
                    for element_dict in relationship_dict["elements"]
                ]
                relationship_diff["peers"] = peer_diffs
            element_diffs.append(relationship_diff)
    return NodeDiff(
        branch=branch_name,
        kind=str(node_dict.get("kind")),
        id=str(node_dict.get("uuid")),
        action=str(node_dict.get("status")),
        display_label=str(node_dict.get("label")),
        elements=element_diffs,
    )


def get_diff_tree_query() -> Query:
    node_structure = {
        "uuid": None,
        "kind": None,
        "status": None,
        "label": None,
        "num_added": None,
        "num_updated": None,
        "num_removed": None,
        "attributes": {
            "name": None,
            "status": None,
            "num_added": None,
            "num_updated": None,
            "num_removed": None,
        },
        "relationships": {
            "name": None,
            "status": None,
            "cardinality": None,
            "num_added": None,
            "num_updated": None,
            "num_removed": None,
            "elements": {
                "status": None,
                "num_added": None,
                "num_updated": None,
                "num_removed": None,
            },
        },
    }

    return Query(
        name="GetDiffTree",
        query={
            "DiffTree": {
                "@filters": {
                    "branch": "$branch_name",
                    "name": "$name",
                    "from_time": "$from_time",
                    "to_time": "$to_time",
                },
                "name": None,
                "to_time": None,
                "from_time": None,
                "base_branch": None,
                "diff_branch": None,
                "num_added": None,
                "num_updated": None,
                "num_removed": None,
                "num_conflicts": None,
                "nodes": node_structure,
            },
        },
        variables={
            "branch_name": str,
            "name": str | None,
            "from_time": datetime | None,
            "to_time": datetime | None,
        },
    )
