"""Pathway visualization — render Pathway graphs as Mermaid flowcharts.

Features:

- **Rich labels** — show prompt snippets, tool names on nodes
- **Execution overlay** — color nodes by execution status from a trace

Usage::

    from core.pathways.visualize import pathway_to_mermaid

    mermaid_str = pathway_to_mermaid(pathway)
    print(mermaid_str)

With execution overlay::

    from core.pathways.visualize import pathway_to_mermaid_with_trace

    mermaid_str = pathway_to_mermaid_with_trace(pathway, trace)
"""

from __future__ import annotations

import re
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable

    from core.observability.trace import ExecutionTrace, TraceSpan
    from core.pathways.pathway import Connection, Pathway


# =============================================================================
# HELPERS
# =============================================================================

_MERMAID_ID_RE = re.compile(r"[^a-zA-Z0-9_]")


def _sanitize_id(raw: str) -> str:
    """Make a string safe for use as a Mermaid node ID."""
    return _MERMAID_ID_RE.sub("_", raw)


def _escape_label(raw: str) -> str:
    """Escape characters that break Mermaid labels inside double quotes."""
    return raw.replace("\\", "\\\\").replace('"', "#quot;")


def _truncate(text: str, max_len: int = 40) -> str:
    """Truncate text and add ellipsis if needed."""
    text = text.replace("\n", " ").strip()
    return text if len(text) <= max_len else f"{text[: max_len - 1]}…"


# =============================================================================
# NODE SHAPE DISPATCH TABLE
# =============================================================================

_SHAPE_FORMATTERS: dict[str, Callable[[str, str], str]] = {
    "input": lambda nid, lbl: f'{nid}(["{lbl}"])',
    "output": lambda nid, lbl: f'{nid}(["{lbl}"])',
    "llm": lambda nid, lbl: f'{nid}{{{{"{lbl}"}}}}',
    "tool": lambda nid, lbl: f'{nid}[["{lbl}"]]',
    "for_each": lambda nid, lbl: f'{nid}(("{lbl}"))',
    "approval": lambda nid, lbl: f'{nid}((("{lbl}")))',
    "pathway": lambda nid, lbl: f'{nid}[["{lbl}"]]',
}


def _default_shape(nid: str, lbl: str) -> str:
    return f'{nid}["{lbl}"]'


# CSS class per node type for Mermaid styling
_NODE_TYPE_CLASSES: dict[str, str] = {
    "input": "nodeInput",
    "output": "nodeOutput",
    "llm": "nodeLLM",
    "tool": "nodeTool",
    "for_each": "nodeLoop",
    "approval": "nodeApproval",
    "pathway": "nodeTool",
}


# =============================================================================
# RICH LABEL BUILDER
# =============================================================================


def _label_llm(node: object) -> str:
    name = getattr(node, "name", None) or getattr(node, "id", "llm")
    prompt = getattr(node, "prompt", "")
    return f"{name}\n{_truncate(prompt, 35)}"


def _label_tool(node: object) -> str:
    name = getattr(node, "name", None) or getattr(node, "id", "tool")
    tool = getattr(node, "tool", "")
    return f"{name}\n{tool}"


def _label_pathway_node(node: object) -> str:
    name = getattr(node, "name", None) or getattr(node, "id", "sub")
    pid = getattr(node, "pathway_id", "?")
    return f"{name}\nsub: {pid}"


def _label_for_each(node: object) -> str:
    name = getattr(node, "name", None) or getattr(node, "id", "each")
    pid = getattr(node, "body_pathway_id", "?")
    return f"{name}\neach: {pid}"


def _label_default(node: object) -> str:
    return getattr(node, "name", None) or getattr(node, "id", "node")


_LABEL_BUILDERS: dict[str, Callable[[object], str]] = {
    "llm": _label_llm,
    "tool": _label_tool,
    "pathway": _label_pathway_node,
    "for_each": _label_for_each,
}


def _build_label(node: object) -> str:
    """Build a rich label for a node based on its type."""
    node_type = getattr(node, "type", "unknown")
    builder = _LABEL_BUILDERS.get(node_type, _label_default)
    return builder(node)


# =============================================================================
# RENDER NODE
# =============================================================================


def _render_node(
    lines: list[str],
    node_id: str,
    node: object,
    node_classes: dict[str, str],
    *,
    indent: str = "    ",
    trace_spans: dict[str, TraceSpan] | None = None,
) -> None:
    """Render a single node definition line."""
    safe_id = _sanitize_id(node_id)
    node_type = getattr(node, "type", "unknown")
    label = _build_label(node)

    # Append timing from trace if available
    if trace_spans and node_id in trace_spans:
        span = trace_spans[node_id]
        if span.duration_ms is not None:
            from core.observability.trace import _format_gantt_duration

            label += f"\n{_format_gantt_duration(span.duration_ms)}"

    label = _escape_label(label)
    formatter = _SHAPE_FORMATTERS.get(node_type, _default_shape)
    lines.append(f"{indent}{formatter(safe_id, label)}")

    css_class = _NODE_TYPE_CLASSES.get(node_type)
    if css_class:
        node_classes[safe_id] = css_class


# =============================================================================
# RENDER CONNECTION
# =============================================================================


def _render_connection(lines: list[str], conn: Connection) -> None:
    """Append a Mermaid edge for a single Connection."""
    from_id = _sanitize_id(conn.from_node)
    to_id = _sanitize_id(conn.to_node)

    label_parts: list[str] = []

    # Non-default port labels
    if conn.from_output != "output" or conn.to_input != "input":
        label_parts.append(f"{conn.from_output} -> {conn.to_input}")

    # When condition label
    if conn.when:
        label_parts.append(_escape_label(conn.when))

    if label_parts:
        label = " | ".join(label_parts)
        lines.append(f'    {from_id} -->|"{label}"| {to_id}')
    else:
        lines.append(f"    {from_id} --> {to_id}")


# =============================================================================
# PUBLIC API: pathway_to_mermaid
# =============================================================================


def pathway_to_mermaid(pathway: Pathway, *, direction: str = "TD") -> str:
    """Render a Pathway as a Mermaid flowchart string.

    Features rich node labels, subgraph grouping for loops, and
    conditional branch annotations.

    Parameters
    ----------
    pathway
        The pathway to visualize.
    direction
        Graph direction. ``"TD"`` (top-down) or ``"LR"`` (left-right).

    Returns
    -------
    str
        A complete Mermaid flowchart definition.
    """
    return _render_pathway_mermaid(pathway, direction=direction, trace_spans=None)


def pathway_to_mermaid_with_trace(
    pathway: Pathway,
    trace: ExecutionTrace,
    *,
    direction: str = "TD",
) -> str:
    """Render a Pathway as a Mermaid flowchart with execution overlay.

    Nodes are colored by execution status (success/failure/skipped/not-run)
    and annotated with timing from the trace.

    Parameters
    ----------
    pathway
        The pathway to visualize.
    trace
        The execution trace to overlay.
    direction
        Graph direction.

    Returns
    -------
    str
        A Mermaid flowchart with execution styling.
    """
    spans_by_id = {s.node_id: s for s in trace.spans}
    return _render_pathway_mermaid(pathway, direction=direction, trace_spans=spans_by_id)


# =============================================================================
# INTERNAL RENDERER
# =============================================================================


def _render_pathway_mermaid(
    pathway: Pathway,
    *,
    direction: str = "TD",
    trace_spans: dict[str, TraceSpan] | None = None,
) -> str:
    """Core rendering logic shared by public API functions."""
    lines: list[str] = [f"flowchart {direction}"]
    node_classes: dict[str, str] = {}

    # -- Node definitions ------------------------------------------------------
    for node_id, node in pathway.nodes.items():
        _render_node(lines, node_id, node, node_classes, trace_spans=trace_spans)

    # -- Connections -----------------------------------------------------------
    for conn in pathway.connections:
        _render_connection(lines, conn)

    # -- Class assignments ----------------------------------------------------
    if trace_spans is not None:
        # Execution overlay: assign status-based classes
        lines.append("")
        for node_id in pathway.nodes:
            safe_id = _sanitize_id(node_id)
            span = trace_spans.get(node_id)
            if span is None:
                node_classes[safe_id] = "notrun"
            elif span.status == "success":
                node_classes[safe_id] = "ok"
            elif span.status == "failure":
                node_classes[safe_id] = "fail"
            else:
                node_classes[safe_id] = "skip"

    classes_by_type: dict[str, list[str]] = {}
    for nid, cls in node_classes.items():
        classes_by_type.setdefault(cls, []).append(nid)

    if classes_by_type:
        lines.append("")
        for cls, nids in classes_by_type.items():
            lines.append(f"    class {','.join(nids)} {cls}")

    return "\n".join(lines) + "\n"


__all__ = [
    "pathway_to_mermaid",
    "pathway_to_mermaid_with_trace",
]
