from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple

import yaml


@dataclass(frozen=True)
class Edge:
    src: str
    dst: str
    label: str


@dataclass(frozen=True)
class Node:
    node_id: str
    label: str
    kind: str  # state | pseudo | event
    shape: str = "rect"  # rect | rounded | pseudo


def _format_guard_value(v: Any) -> str:
    if v is True:
        return "true"
    if v is False:
        return "false"
    if v is None:
        return "null"
    if isinstance(v, (int, float)):
        return str(v)
    if isinstance(v, str):
        return v
    if isinstance(v, list):
        # Keep it readable; lists are rare in guards.
        return "[" + ",".join(_format_guard_value(x) for x in v[:6]) + ("]" if len(v) <= 6 else ",…]")
    return str(v)


def _guard_summary(when: Mapping[str, Any], *, max_len: int = 42) -> str:
    parts: List[str] = []
    for k in sorted(when.keys()):
        parts.append(f"{k}={_format_guard_value(when[k])}")
    s = " ".join(parts)
    if len(s) > max_len:
        return s[:max_len] + "…"
    return s


def _shorten(s: str, max_len: int) -> str:
    s = s.strip().replace("\n", " ")
    if len(s) <= max_len:
        return s
    return s[:max_len] + "…"


def _desc_prefix(desc: str, *, max_len: int) -> str:
    """Pick a high-signal prefix for edge labels.

    We prefer the text before common separators so the main meaning appears first.
    """
    s = desc.strip().replace("\n", " ")
    for sep in ("：", "→", "->", "（", "(", "，", ",", ";", "。"): 
        if sep in s:
            head = s.split(sep, 1)[0].strip()
            if 4 <= len(head) <= max_len:
                return head
    return _shorten(s, max_len)


def _branch_suffix(i: int) -> str:
    # a, b, c ...; if too many, fall back to xN.
    if 0 <= i < 26:
        return chr(ord("a") + i)
    return f"x{i + 1}"


def _event_outcome_tag(event_id: str, when: Mapping[str, Any]) -> str:
    """Optional human-friendly tag appended to an edge label.

    Keep it minimal and only for high-value cases to avoid clutter.
    """
    if event_id == "E16":
        # Auto-bind: show the effective matching shape (only the two requested ones).
        outcome = when.get("auto_bind_outcome")
        if outcome == "ONE_TO_ONE":
            return "（1:1）"
        if outcome == "ZERO":
            return "（N:0）"
    return ""


def _load_yaml(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    if not isinstance(data, dict):
        raise TypeError("Top-level YAML must be a mapping")
    return data


def _get_state_status(states: Mapping[str, Any], state_id: str) -> Optional[str]:
    s = states.get(state_id)
    if not isinstance(s, dict):
        return None
    status = s.get("status")
    return status if isinstance(status, str) else None


def _iter_to_map_targets(states: Mapping[str, Any], from_state: str, to_map: Mapping[str, Any]) -> Set[str]:
    key = to_map.get("key")
    mp = to_map.get("map")
    if key != "self.status" or not isinstance(mp, dict):
        # Unsupported mapping; fall back to all targets
        return {v for v in mp.values()} if isinstance(mp, dict) else set()

    status = _get_state_status(states, from_state)
    if status and status != "*" and status in mp and isinstance(mp[status], str):
        return {mp[status]}

    # Wildcard or unknown status: return all possible targets
    targets: Set[str] = set()
    for v in mp.values():
        if isinstance(v, str):
            targets.add(v)
    return targets


def _collect_state_edges(cfg: Mapping[str, Any]) -> List[Edge]:
    states = cfg.get("states") if isinstance(cfg.get("states"), dict) else {}
    transitions = cfg.get("transitions") if isinstance(cfg.get("transitions"), dict) else {}

    # Aggregate multiple transitions between the same two states.
    edge_events: Dict[Tuple[str, str], Set[str]] = {}

    for tid, tdef in transitions.items():
        if not isinstance(tdef, dict):
            continue
        frm = tdef.get("from")
        if not isinstance(frm, list) or not frm:
            continue

        outcomes = tdef.get("outcomes")
        if not isinstance(outcomes, list) or not outcomes:
            continue

        tid_str = str(tid)
        desc = tdef.get("desc")
        desc_short = _desc_prefix(desc, max_len=16) if isinstance(desc, str) and desc else ""

        for src in frm:
            if not isinstance(src, str):
                continue
            if src == "*":
                # For visualization, we don't expand "*" because it explodes.
                continue

            src_action: Optional[str] = None
            src_state = states.get(src)
            if isinstance(src_state, dict):
                action_val = src_state.get("action")
                if isinstance(action_val, str):
                    src_action = action_val

            for oi, o in enumerate(outcomes):
                if not isinstance(o, dict):
                    continue

                when = o.get("when")
                if not isinstance(when, dict):
                    when = {}

                # Respect execute_action guard in graph generation so we don't
                # draw branches that are impossible for the source state's action.
                if src_action is not None and "execute_action" in when:
                    guarded_action = when.get("execute_action")
                    if isinstance(guarded_action, str) and src_action != guarded_action:
                        continue
                    if isinstance(guarded_action, list) and all(isinstance(v, str) for v in guarded_action):
                        if src_action not in guarded_action:
                            continue

                targets: Set[str] = set()
                if "to" in o and isinstance(o.get("to"), str):
                    targets.add(o["to"])
                elif "to_map" in o and isinstance(o.get("to_map"), dict):
                    tm = o["to_map"]
                    if tm.get("key") == "self.status" and isinstance(tm.get("map"), dict):
                        for status_key, dst_state in tm["map"].items():
                            if isinstance(dst_state, str) and isinstance(status_key, str):
                                targets.add(dst_state)
                        continue
                    targets |= _iter_to_map_targets(states, src, tm)

                for dst in sorted(targets):
                    label = f"{tid_str}{_branch_suffix(oi)}"
                    if desc_short:
                        tag = _event_outcome_tag(tid_str, when)
                        label = f"{label} {desc_short}{tag}"
                    edge_events.setdefault((src, dst), set()).add(label)

    edges: List[Edge] = []
    for (src, dst), tids in sorted(edge_events.items()):
        label = ", ".join(sorted(tids))
        edges.append(Edge(src=src, dst=dst, label=label))
    return edges


def _iter_state_nodes(cfg: Mapping[str, Any]) -> List[Node]:
    nodes: List[Node] = []
    states = cfg.get("states") if isinstance(cfg.get("states"), dict) else {}
    pseudo = cfg.get("pseudo_states") if isinstance(cfg.get("pseudo_states"), dict) else {}

    for sid, sdef in states.items():
        if not isinstance(sdef, dict):
            continue
        desc = sdef.get("desc")
        label = sid
        if isinstance(desc, str) and desc:
            label = f"{sid}\\n{desc}"
        shape = "rect"
        if sdef.get("shape") == "rounded":
            shape = "rounded"
        elif sdef.get("shape") == "pseudo":
            shape = "pseudo"
        nodes.append(Node(node_id=sid, label=label, kind="state", shape=shape))

    for pid, pdef in pseudo.items():
        label = pid
        if isinstance(pdef, dict) and isinstance(pdef.get("desc"), str) and pdef.get("desc"):
            label = f"{pid}\\n{pdef['desc']}"
        nodes.append(Node(node_id=pid, label=label, kind="pseudo", shape="pseudo"))

    return nodes


def render_dot(cfg: Mapping[str, Any]) -> str:
    nodes = _iter_state_nodes(cfg)
    edges = _collect_state_edges(cfg)
    node_ids = {n.node_id for n in nodes}

    lines: List[str] = []
    lines.append("digraph NodeStateMachine {")
    lines.append("  rankdir=LR;")
    lines.append("  node [fontname=\"Consolas\"];\n")

    for n in nodes:
        safe_label = n.label.replace('"', "\\\"")
        shape = "box"
        attrs: List[str] = [f"label=\"{safe_label}\"", f"shape={shape}"]
        if n.kind == "pseudo":
            attrs = [f"label=\"{safe_label}\"", "shape=oval"]
        elif n.shape == "pseudo":
            attrs = [f"label=\"{safe_label}\"", "shape=oval"]
        elif n.shape == "rounded":
            attrs.append("style=rounded")
        lines.append(f"  \"{n.node_id}\" [{', '.join(attrs)}]; ")

    lines.append("")
    for e in edges:
        safe_label = e.label.replace('"', "\\\"")
        if safe_label:
            lines.append(f"  \"{e.src}\" -> \"{e.dst}\" [label=\"{safe_label}\"]; ")
        else:
            lines.append(f"  \"{e.src}\" -> \"{e.dst}\"; ")

    if {"S13", "S06"}.issubset(node_ids):
        lines.append(
            '  "S13" -> "S06" [style=dashed, label="副作用: E16识别需创建后，对侧spawn S06"]; '
        )

    lines.append("}")
    return "\n".join(lines)


def render_mermaid(cfg: Mapping[str, Any]) -> str:
    # Mermaid flowchart is more forgiving than stateDiagram-v2 when labels are long.
    # Subgraphs are always enabled to reflect the five-stage mental model.
    nodes = _iter_state_nodes(cfg)
    edges = _collect_state_edges(cfg)
    node_ids = {n.node_id for n in nodes}

    def _group(node_id: str) -> str:
        if node_id in {"S16", "S15", "S17", "S00"}:
            return "BOOT"
        if node_id in {"S01", "S02", "S03", "S04", "S05"}:
            return "BIND"
        if node_id in {"S06", "S10C", "S11C", "S12C"}:
            return "CREATE"
        if node_id in {"S07", "S08", "S10U", "S11U", "S12U", "S14"}:
            return "UPDATE"
        if node_id in {"S09", "S10D", "S11D", "S12D"}:
            return "DELETE"
        return "BIND"

    group_titles = {
        "BOOT": "Bootstrap / Reset",
        "BIND": "Binding",
        "CREATE": "Create",
        "UPDATE": "Update",
        "DELETE": "Delete",
    }

    grouped_nodes: Dict[str, List[Node]] = {k: [] for k in group_titles.keys()}
    for n in nodes:
        grouped_nodes[_group(n.node_id)].append(n)

    lines: List[str] = []
    lines.append("flowchart LR")

    for group_id in ["BOOT", "BIND", "CREATE", "UPDATE", "DELETE"]:
        lines.append("")
        lines.append(f"  subgraph {group_id}[\"{group_titles[group_id]}\"]")
        for n in sorted(grouped_nodes[group_id], key=lambda x: x.node_id):
            safe = n.label.replace("\\n", "<br/>").replace("\n", "<br/>").replace("\"", "'")
            if n.kind == "pseudo" or n.shape == "pseudo":
                lines.append(f"    {n.node_id}([\"{safe}\"]) ")
            elif n.shape == "rounded":
                lines.append(f"    {n.node_id}(\"{safe}\")")
            else:
                lines.append(f"    {n.node_id}[\"{safe}\"]")
        lines.append("  end")

    lines.append("")
    for e in edges:
        safe_label = e.label.replace("\"", "'")
        if safe_label:
            lines.append(f"  {e.src} -- \"{safe_label}\" --> {e.dst}")
        else:
            lines.append(f"  {e.src} --> {e.dst}")

    if {"S13", "S06"}.issubset(node_ids):
        lines.append('  S13 -. "副作用: E16识别需创建后，对侧spawn S06" .-> S06')

    return "\n".join(lines)


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Render state machine graph from YAML")
    parser.add_argument(
        "--config",
        default=str(Path(__file__).resolve().parents[1] / "sync_state_machine" / "config" / "node_state_machine.yaml"),
        help="Path to YAML config",
    )
    parser.add_argument(
        "--format",
        choices=["dot", "mermaid"],
        default="mermaid",
        help="Output format",
    )
    parser.add_argument(
        "--out",
        default=str(Path(__file__).resolve().parents[1] / "docs" / "state_machine.mmd"),
        help="Output file path, or '-' for stdout",
    )

    args = parser.parse_args(argv)

    cfg = _load_yaml(Path(args.config))
    if args.format == "dot":
        out = render_dot(cfg)
    else:
        out = render_mermaid(cfg)

    if args.out == "-":
        print(out)
    else:
        Path(args.out).write_text(out, encoding="utf-8")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())