"""
mermaid_graph.schema
=================
Dataclass definitions that mirror the node_link_data-extended JSON schema
for Mermaid diagrams.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .exceptions import MermaidValidationError


# ── Shape / Line / Arrow enums as string constants ──────────────────────

SHAPES = {
    "rect", "round_rect", "stadium", "diamond", "hexagon", "circle",
    "subroutine", "cylinder", "asymmetric", "parallelogram",
    "parallelogram_alt", "trapezoid", "trapezoid_alt", "double_circle",
}

LINE_TYPES = {"solid", "dotted", "thick"}

ARROW_TYPES = {"normal", "none", "circle", "cross"}

DIRECTIONS = {"TD", "TB", "BT", "LR", "RL"}


# ── Core dataclasses ────────────────────────────────────────────────────

@dataclass
class MermaidNode:
    id: str
    label: str | None = None
    shape: str = "rect"
    css_class: str | None = None
    style: dict[str, str] | None = None

    def __post_init__(self) -> None:
        if self.shape not in SHAPES:
            raise MermaidValidationError(
                f"Unknown shape {self.shape!r}. Valid shapes: {sorted(SHAPES)}"
            )

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"id": self.id}
        if self.label is not None:
            d["label"] = self.label
        d["shape"] = self.shape
        if self.css_class is not None:
            d["css_class"] = self.css_class
        if self.style:
            d["style"] = self.style
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> MermaidNode:
        return cls(
            id=d["id"],
            label=d.get("label"),
            shape=d.get("shape", "rect"),
            css_class=d.get("css_class"),
            style=d.get("style"),
        )


@dataclass
class MermaidLink:
    source: str
    target: str
    label: str | None = None
    line_type: str = "solid"
    arrow: str = "normal"
    link_index: int | None = None
    style: dict[str, str] | None = None

    def __post_init__(self) -> None:
        if self.line_type not in LINE_TYPES:
            raise MermaidValidationError(
                f"Unknown line_type {self.line_type!r}. "
                f"Valid types: {sorted(LINE_TYPES)}"
            )
        if self.arrow not in ARROW_TYPES:
            raise MermaidValidationError(
                f"Unknown arrow {self.arrow!r}. "
                f"Valid types: {sorted(ARROW_TYPES)}"
            )

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "source": self.source,
            "target": self.target,
        }
        if self.label is not None:
            d["label"] = self.label
        d["line_type"] = self.line_type
        d["arrow"] = self.arrow
        if self.link_index is not None:
            d["link_index"] = self.link_index
        if self.style:
            d["style"] = self.style
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> MermaidLink:
        return cls(
            source=d["source"],
            target=d["target"],
            label=d.get("label"),
            line_type=d.get("line_type", "solid"),
            arrow=d.get("arrow", "normal"),
            link_index=d.get("link_index"),
            style=d.get("style"),
        )


@dataclass
class Subgraph:
    id: str
    label: str | None = None
    nodes: list[str] = field(default_factory=list)
    subgraphs: list[Subgraph] = field(default_factory=list)
    direction: str | None = None
    style: dict[str, str] | None = None

    def __post_init__(self) -> None:
        if self.direction is not None and self.direction not in DIRECTIONS:
            raise MermaidValidationError(
                f"Unknown direction {self.direction!r}. "
                f"Valid directions: {sorted(DIRECTIONS)}"
            )

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"id": self.id}
        if self.label is not None:
            d["label"] = self.label
        d["nodes"] = self.nodes
        d["subgraphs"] = [sg.to_dict() for sg in self.subgraphs]
        if self.direction is not None:
            d["direction"] = self.direction
        if self.style:
            d["style"] = self.style
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Subgraph:
        return cls(
            id=d["id"],
            label=d.get("label"),
            nodes=d.get("nodes", []),
            subgraphs=[cls.from_dict(sg) for sg in d.get("subgraphs", [])],
            direction=d.get("direction"),
            style=d.get("style"),
        )


@dataclass
class MermaidGraph:
    """Top-level container matching the extended node_link_data format."""

    directed: bool = True
    multigraph: bool = False
    diagram_type: str = "flowchart"
    direction: str = "TD"
    title: str | None = None
    class_defs: dict[str, dict[str, str]] = field(default_factory=dict)
    subgraphs: list[Subgraph] = field(default_factory=list)
    link_styles: dict[str, dict[str, str]] = field(default_factory=dict)
    nodes: list[MermaidNode] = field(default_factory=list)
    links: list[MermaidLink] = field(default_factory=list)

    def __post_init__(self) -> None:
        if self.direction not in DIRECTIONS:
            raise MermaidValidationError(
                f"Unknown direction {self.direction!r}. "
                f"Valid directions: {sorted(DIRECTIONS)}"
            )

    # ── Serialisation helpers ───────────────────────────────────────

    def to_dict(self) -> dict[str, Any]:
        """Return the extended node_link_data dict."""
        graph_meta: dict[str, Any] = {
            "diagram_type": self.diagram_type,
            "direction": self.direction,
        }
        if self.title:
            graph_meta["title"] = self.title
        if self.class_defs:
            graph_meta["class_defs"] = self.class_defs
        if self.subgraphs:
            graph_meta["subgraphs"] = [sg.to_dict() for sg in self.subgraphs]
        if self.link_styles:
            graph_meta["link_styles"] = {
                str(k): v for k, v in self.link_styles.items()
            }

        return {
            "directed": self.directed,
            "multigraph": self.multigraph,
            "graph": graph_meta,
            "nodes": [n.to_dict() for n in self.nodes],
            "links": [lk.to_dict() for lk in self.links],
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> MermaidGraph:
        g = d.get("graph", {})
        return cls(
            directed=d.get("directed", True),
            multigraph=d.get("multigraph", False),
            diagram_type=g.get("diagram_type", "flowchart"),
            direction=g.get("direction", "TD"),
            title=g.get("title"),
            class_defs=g.get("class_defs", {}),
            subgraphs=[Subgraph.from_dict(sg) for sg in g.get("subgraphs", [])],
            link_styles={
                str(k): v for k, v in g.get("link_styles", {}).items()
            },
            nodes=[MermaidNode.from_dict(n) for n in d.get("nodes", [])],
            links=[MermaidLink.from_dict(lk) for lk in d.get("links", [])],
        )
