"""Typed workspace graph for the cognitive runtime."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any, Iterable
from uuid import uuid4


class NodeType(StrEnum):
    TASK = "task"
    GOAL = "goal"
    SUBGOAL = "subgoal"
    HYPOTHESIS = "hypothesis"
    EVIDENCE = "evidence"
    CONSTRAINT = "constraint"
    ACTION = "action"
    RESULT = "result"
    MEMORY = "memory"
    SKILL = "skill"
    ANSWER = "answer"


@dataclass(slots=True)
class Node:
    type: NodeType
    label: str
    data: dict[str, Any] = field(default_factory=dict)
    confidence: float = 1.0
    node_id: str = field(default_factory=lambda: uuid4().hex)

    def __post_init__(self) -> None:
        if not isinstance(self.type, NodeType):
            self.type = NodeType(str(self.type))
        if not 0.0 <= float(self.confidence) <= 1.0:
            raise ValueError("confidence must be between 0 and 1")

    def to_dict(self) -> dict[str, Any]:
        return {
            "node_id": self.node_id,
            "type": self.type.value,
            "label": self.label,
            "data": self.data,
            "confidence": self.confidence,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Node":
        return cls(
            node_id=str(data["node_id"]),
            type=NodeType(str(data["type"])),
            label=str(data["label"]),
            data=dict(data.get("data") or {}),
            confidence=float(data.get("confidence", 1.0)),
        )


@dataclass(slots=True)
class Edge:
    source: str
    target: str
    relation: str
    weight: float = 1.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "source": self.source,
            "target": self.target,
            "relation": self.relation,
            "weight": self.weight,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Edge":
        return cls(
            source=str(data["source"]),
            target=str(data["target"]),
            relation=str(data["relation"]),
            weight=float(data.get("weight", 1.0)),
        )


@dataclass(slots=True)
class Workspace:
    task: str
    nodes: dict[str, Node] = field(default_factory=dict)
    edges: list[Edge] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_task(cls, task: str) -> "Workspace":
        workspace = cls(task=task, metadata={"stage": "created", "step": 0})
        workspace.add_node(NodeType.TASK, task, data={"raw": task})
        return workspace

    def add_node(
        self,
        node_type: NodeType | str,
        label: str,
        data: dict[str, Any] | None = None,
        confidence: float = 1.0,
    ) -> Node:
        node = Node(NodeType(str(node_type)), label, data or {}, confidence)
        self.nodes[node.node_id] = node
        return node

    def add_edge(self, source: str | Node, target: str | Node, relation: str, weight: float = 1.0) -> Edge:
        source_id = source.node_id if isinstance(source, Node) else str(source)
        target_id = target.node_id if isinstance(target, Node) else str(target)
        if source_id not in self.nodes:
            raise KeyError(f"Unknown source node: {source_id}")
        if target_id not in self.nodes:
            raise KeyError(f"Unknown target node: {target_id}")
        edge = Edge(source=source_id, target=target_id, relation=relation, weight=weight)
        self.edges.append(edge)
        return edge

    def find(self, node_type: NodeType | str) -> list[Node]:
        expected = NodeType(str(node_type))
        return [node for node in self.nodes.values() if node.type == expected]

    def labels(self, node_type: NodeType | str) -> list[str]:
        return [node.label for node in self.find(node_type)]

    def has_answer(self) -> bool:
        return bool(self.find(NodeType.ANSWER))

    def increment_step(self) -> int:
        step = int(self.metadata.get("step", 0)) + 1
        self.metadata["step"] = step
        return step

    def set_stage(self, stage: str) -> None:
        self.metadata["stage"] = stage

    def to_dict(self) -> dict[str, Any]:
        return {
            "task": self.task,
            "nodes": [node.to_dict() for node in self.nodes.values()],
            "edges": [edge.to_dict() for edge in self.edges],
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Workspace":
        workspace = cls(task=str(data["task"]), metadata=dict(data.get("metadata") or {}))
        for node_data in data.get("nodes", []):
            node = Node.from_dict(node_data)
            workspace.nodes[node.node_id] = node
        workspace.edges = [Edge.from_dict(edge) for edge in data.get("edges", [])]
        return workspace

    def summary(self) -> str:
        counts: dict[str, int] = {}
        for node in self.nodes.values():
            counts[node.type.value] = counts.get(node.type.value, 0) + 1
        parts = [f"{key}={value}" for key, value in sorted(counts.items())]
        return ", ".join(parts)

    def iter_nodes(self, types: Iterable[NodeType] | None = None) -> Iterable[Node]:
        allowed = set(types) if types is not None else None
        for node in self.nodes.values():
            if allowed is None or node.type in allowed:
                yield node
