from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Callable, Generic, Mapping, Protocol, TypeVar

from .architectures import AcceptancePolicy, Categorizer, Critic, Generator


StateT = TypeVar("StateT")
CandidateT = TypeVar("CandidateT")
CritiqueT = TypeVar("CritiqueT")
CategoryT = TypeVar("CategoryT")


@dataclass(frozen=True)
class GraphValue:
    value: object
    kind: str


@dataclass(frozen=True)
class GraphNodeResult:
    outputs: dict[str, GraphValue]


@dataclass(frozen=True)
class GraphNodeSpec:
    node_id: str
    node_type: str
    input_names: tuple[str, ...]
    output_names: tuple[str, ...]


@dataclass(frozen=True)
class GraphEdgeSpec:
    source_node_id: str
    source_output: str
    target_node_id: str
    target_input: str


@dataclass(frozen=True)
class GraphSchema:
    nodes: tuple[GraphNodeSpec, ...]
    edges: tuple[GraphEdgeSpec, ...]

    def to_dict(self) -> dict[str, object]:
        return asdict(self)

    def save_json(self, path: str | Path) -> None:
        from .artifacts import save_artifact_json

        save_artifact_json("graph_schema", self, path)


class GraphNode(Protocol):
    node_id: str
    node_type: str
    input_names: tuple[str, ...]
    output_names: tuple[str, ...]

    def run(self, inputs: Mapping[str, object]) -> GraphNodeResult:
        ...

    def spec(self) -> GraphNodeSpec:
        ...


class FunctionalNode:
    def __init__(
        self,
        *,
        node_id: str,
        node_type: str,
        input_names: tuple[str, ...],
        output_names: tuple[str, ...],
        fn: Callable[[Mapping[str, object]], GraphNodeResult],
    ) -> None:
        self.node_id = node_id
        self.node_type = node_type
        self.input_names = input_names
        self.output_names = output_names
        self._fn = fn

    def run(self, inputs: Mapping[str, object]) -> GraphNodeResult:
        return self._fn(inputs)

    def spec(self) -> GraphNodeSpec:
        return GraphNodeSpec(
            node_id=self.node_id,
            node_type=self.node_type,
            input_names=self.input_names,
            output_names=self.output_names,
        )


def generator_node(
    node_id: str,
    generator: Generator[StateT, CandidateT],
) -> FunctionalNode:
    def run(inputs: Mapping[str, object]) -> GraphNodeResult:
        state = inputs["state"]
        candidate = generator.generate(state)  # type: ignore[arg-type]
        return GraphNodeResult(
            outputs={
                "candidate": GraphValue(candidate, "candidate"),
            }
        )

    return FunctionalNode(
        node_id=node_id,
        node_type="generator",
        input_names=("state",),
        output_names=("candidate",),
        fn=run,
    )


def critic_node(
    node_id: str,
    critic: Critic[StateT, CandidateT, CritiqueT],
) -> FunctionalNode:
    def run(inputs: Mapping[str, object]) -> GraphNodeResult:
        critique = critic.critique(
            inputs["state"],  # type: ignore[arg-type]
            inputs["candidate"],  # type: ignore[arg-type]
        )
        return GraphNodeResult(
            outputs={
                "critique": GraphValue(critique, "critique"),
            }
        )

    return FunctionalNode(
        node_id=node_id,
        node_type="critic",
        input_names=("state", "candidate"),
        output_names=("critique",),
        fn=run,
    )


def categorizer_node(
    node_id: str,
    categorizer: Categorizer[StateT, CandidateT, CategoryT],
) -> FunctionalNode:
    def run(inputs: Mapping[str, object]) -> GraphNodeResult:
        category = categorizer.categorize(
            inputs["state"],  # type: ignore[arg-type]
            inputs["candidate"],  # type: ignore[arg-type]
        )
        return GraphNodeResult(
            outputs={
                "category": GraphValue(category, "category"),
            }
        )

    return FunctionalNode(
        node_id=node_id,
        node_type="categorizer",
        input_names=("state", "candidate"),
        output_names=("category",),
        fn=run,
    )


def policy_node(
    node_id: str,
    policy: AcceptancePolicy[StateT, CandidateT, CritiqueT, CategoryT],
) -> FunctionalNode:
    def run(inputs: Mapping[str, object]) -> GraphNodeResult:
        decision = policy.decide(
            inputs["state"],  # type: ignore[arg-type]
            inputs["candidate"],  # type: ignore[arg-type]
            inputs["critique"],  # type: ignore[arg-type]
            inputs["category"],  # type: ignore[arg-type]
        )
        return GraphNodeResult(
            outputs={
                "decision": GraphValue(decision, "decision"),
                "accepted": GraphValue(decision.accepted, "boolean"),
            }
        )

    return FunctionalNode(
        node_id=node_id,
        node_type="policy",
        input_names=("state", "candidate", "critique", "category"),
        output_names=("decision", "accepted"),
        fn=run,
    )
