"""Workflow visualization helpers (Mermaid + optional Graphviz DOT)."""

from __future__ import annotations

import re
from pathlib import Path
from typing import Any, Literal


class WorkflowVisualizer:
    """Render Agentbyte workflow graphs into text-based visualization formats."""

    @classmethod
    def to_mermaid(
        cls,
        workflow: Any,
        view: Literal["topology", "detailed"] = "topology",
        include_validation: bool = False,
    ) -> str:
        """Return Mermaid flowchart markup for a workflow."""
        lines: list[str] = ["flowchart TD"]

        validation_warnings: list[str] = []
        if include_validation and hasattr(workflow, "validate_workflow"):
            try:
                validation_result = workflow.validate_workflow()
                validation_warnings = list(getattr(validation_result, "warnings", []) or [])
            except Exception:
                validation_warnings = []

        if view == "detailed":
            lines.extend(
                [
                    "    classDef start fill:#d1fae5,stroke:#10b981,stroke-width:2px",
                    "    classDef endStep fill:#fee2e2,stroke:#ef4444,stroke-width:2px",
                    "    classDef agent fill:#dbeafe,stroke:#3b82f6,stroke-width:2px",
                    "    classDef tool fill:#f3e8ff,stroke:#8b5cf6,stroke-dasharray: 4 2",
                    "    classDef step fill:#f8fafc,stroke:#64748b",
                ]
            )

        for step_id, step in cls._sorted_steps(workflow):
            node_id = cls._safe_node_id(step_id)
            node_label = cls._node_label(step_id, step, workflow, detailed=view == "detailed")
            lines.append(f'    {node_id}["{node_label}"]')

            if view == "detailed":
                lines.extend(cls._mermaid_agent_and_tools(node_id, step))
                css_classes = ["step"]
                if step_id == getattr(workflow, "start_step_id", None):
                    css_classes.append("start")
                if step_id in (getattr(workflow, "end_step_ids", []) or []):
                    css_classes.append("endStep")
                lines.append(f"    class {node_id} {','.join(css_classes)}")

        for edge in cls._sorted_edges(workflow):
            source = cls._safe_node_id(cls._edge_source(edge))
            target = cls._safe_node_id(cls._edge_target(edge))
            label = cls._condition_label(getattr(edge, "condition", None))
            lines.append(f'    {source} -->|"{label}"| {target}')

        if validation_warnings:
            lines.append("    %% Validation warnings")
            for warning in validation_warnings:
                lines.append(f"    %% {warning}")

            if view == "detailed":
                warning_nodes = cls._mermaid_warning_nodes(validation_warnings)
                if warning_nodes:
                    lines.extend(warning_nodes)

        return "\n".join(lines)

    @classmethod
    def to_dot(
        cls,
        workflow: Any,
        view: Literal["topology", "detailed"] = "topology",
    ) -> str:
        """Return Graphviz DOT source for a workflow.

        Raises RuntimeError when graphviz is not installed.
        """
        try:
            from graphviz import Digraph
        except ImportError as exc:
            raise RuntimeError(
                "Graphviz is optional and not installed. Install with `uv sync --group viz`."
            ) from exc

        graph = Digraph(name=getattr(getattr(workflow, "config", None), "name", "workflow"))
        graph.attr(rankdir="TB")

        for step_id, step in cls._sorted_steps(workflow):
            node_id = cls._safe_node_id(step_id)
            graph.node(node_id, cls._node_label(step_id, step, workflow, detailed=view == "detailed").replace("<br/>", "\\n"))

            if view == "detailed":
                agent = getattr(step, "agent", None)
                if agent is None:
                    continue
                agent_name = getattr(agent, "name", "agent")
                agent_node = f"{node_id}_agent"
                graph.node(agent_node, f"agent: {agent_name}", shape="ellipse")
                graph.edge(node_id, agent_node)
                for index, tool in enumerate(getattr(agent, "tools", []) or []):
                    tool_name = getattr(tool, "name", f"tool_{index}")
                    tool_node = f"{node_id}_tool_{index}"
                    graph.node(tool_node, f"tool: {tool_name}", shape="box")
                    graph.edge(agent_node, tool_node, style="dashed", label="uses")

        for edge in cls._sorted_edges(workflow):
            graph.edge(
                cls._safe_node_id(cls._edge_source(edge)),
                cls._safe_node_id(cls._edge_target(edge)),
                label=cls._condition_label(getattr(edge, "condition", None)),
            )

        return graph.source

    @staticmethod
    def save(content: str, out_path: str | Path) -> Path:
        """Write graph markup to disk and return the output path."""
        path = Path(out_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8")
        return path

    @staticmethod
    def _sorted_steps(workflow: Any) -> list[tuple[str, Any]]:
        steps = getattr(workflow, "steps", {}) or {}
        return sorted(steps.items(), key=lambda item: item[0])

    @staticmethod
    def _sorted_edges(workflow: Any) -> list[Any]:
        edges = getattr(workflow, "edges", []) or []
        return sorted(edges, key=lambda edge: (WorkflowVisualizer._edge_source(edge), WorkflowVisualizer._edge_target(edge)))

    @staticmethod
    def _edge_source(edge: Any) -> str:
        return getattr(edge, "from_step", getattr(edge, "source", ""))

    @staticmethod
    def _edge_target(edge: Any) -> str:
        return getattr(edge, "to_step", getattr(edge, "target", ""))

    @staticmethod
    def _safe_node_id(value: str) -> str:
        chars = [char if char.isalnum() else "_" for char in value]
        return "".join(chars)

    @staticmethod
    def _condition_label(condition: Any) -> str:
        if condition is None:
            return "always"
        condition_type = getattr(condition, "type", "always")
        if condition_type == "always":
            return "always"
        field = getattr(condition, "field", None)
        operator = getattr(condition, "operator", None)
        value = getattr(condition, "value", None)
        if field and operator:
            return f"{condition_type}: {field} {operator} {value}"
        return str(condition_type)

    @staticmethod
    def _node_label(step_id: str, step: Any, workflow: Any, detailed: bool) -> str:
        if not detailed:
            marker = ""
            if step_id == getattr(workflow, "start_step_id", None):
                marker += " [start]"
            if step_id in (getattr(workflow, "end_step_ids", []) or []):
                marker += " [end]"
            return f"{step_id}{marker}"

        label_parts = [step_id, f"type={step.__class__.__name__}"]
        if getattr(step, "metadata", None) is not None:
            label_parts.append(f"name={step.metadata.name}")
        if step_id == getattr(workflow, "start_step_id", None):
            label_parts.append("START")
        if step_id in (getattr(workflow, "end_step_ids", []) or []):
            label_parts.append("END")
        return "<br/>".join(label_parts)

    @classmethod
    def _mermaid_warning_nodes(cls, warnings: list[str]) -> list[str]:
        lines: list[str] = [
            "    classDef warning fill:#fef3c7,stroke:#f59e0b,stroke-dasharray: 3 3",
        ]
        count = 0

        for warning in warnings:
            step_id = cls._extract_step_id_from_warning(warning)
            if not step_id:
                continue

            safe_step_id = cls._safe_node_id(step_id)
            warning_node_id = f"{safe_step_id}_warning_{count}"
            warning_label = warning.replace('"', "'")
            lines.append(f'    {warning_node_id}["⚠ {warning_label}"]')
            lines.append(f"    {safe_step_id} -.warning.-> {warning_node_id}")
            lines.append(f"    class {warning_node_id} warning")
            count += 1

        return lines if count > 0 else []

    @staticmethod
    def _extract_step_id_from_warning(warning: str) -> str | None:
        step_match = re.search(r"Step '([^']+)'", warning)
        if step_match:
            return step_match.group(1)

        from_match = re.search(r"from '([^']+)'", warning)
        if from_match:
            return from_match.group(1)

        return None

    @staticmethod
    def _mermaid_agent_and_tools(node_id: str, step: Any) -> list[str]:
        lines: list[str] = []
        if not hasattr(step, "agent"):
            return lines

        agent = getattr(step, "agent")
        agent_name = getattr(agent, "name", "agent")
        lines.append(f'    {node_id}_agent(("agent: {agent_name}"))')
        lines.append(f"    {node_id} --> {node_id}_agent")
        lines.append(f"    class {node_id}_agent agent")

        tools = getattr(agent, "tools", []) or []
        for index, tool in enumerate(tools):
            tool_name = getattr(tool, "name", f"tool_{index}")
            tool_node = f"{node_id}_tool_{index}"
            lines.append(f'    {tool_node}[["tool: {tool_name}"]]')
            lines.append(f"    {node_id}_agent -.uses.-> {tool_node}")
            lines.append(f"    class {tool_node} tool")

        return lines


__all__ = ["WorkflowVisualizer"]
