"""Debug and trace handlers for pathway observability events.

All handlers consume ``PathwayEvent`` directly from the ``VMEventEmitter``.
Attach them with ``vm.add_listener(handler)``.
"""

from __future__ import annotations

import json
import logging
import sys
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, ClassVar

if TYPE_CHECKING:
    from collections.abc import Callable

    from core.pathways.models.execution import PathwayEvent

_trace_log = logging.getLogger("albusos.trace")


# =============================================================================
# COLORS
# =============================================================================


class Colors:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"


def _truncate(s: str, max_len: int = 80) -> str:
    return s if len(s) <= max_len else f"{s[: max_len - 3]}..."


def _format_duration(ms: float | None) -> str:
    """Format duration nicely."""
    if ms is None:
        return "—"
    if ms < 1:
        return "<1ms"
    if ms < 1000:
        return f"{ms:.0f}ms"
    if ms < 60000:
        return f"{ms / 1000:.1f}s"
    return f"{ms / 60000:.1f}m"


# =============================================================================
# DEBUG HANDLER
# =============================================================================


class DebugHandler:
    """Pretty-prints ``PathwayEvent`` to stderr for debugging.

    Attach to any ``PathwayVM``::

        handler = DebugHandler()
        vm.add_listener(handler)
    """

    def __init__(self, show_timestamps: bool = True, verbose: bool = False) -> None:
        self.show_timestamps = show_timestamps
        self.verbose = verbose
        self._indent = 0

    def __call__(self, event: PathwayEvent) -> None:
        handler = _DEBUG_DISPATCH.get(event.event_type)
        if handler:
            handler(self, event)
        elif self.verbose:
            self._print(f"{Colors.DIM}[{event.event_type}]{Colors.RESET}")

    def _prefix(self) -> str:
        ts = (
            f"{Colors.DIM}{datetime.now(tz=UTC).strftime('%H:%M:%S.%f')[:-3]}{Colors.RESET} "
            if self.show_timestamps
            else ""
        )
        return f"{ts}{'  ' * self._indent}"

    def _print(self, msg: str) -> None:
        sys.stderr.write(f"{self._prefix()}{msg}\n")
        sys.stderr.flush()

    # -- Per-event handlers ---------------------------------------------------

    def _on_execution_started(self, event: PathwayEvent) -> None:
        pathway_id = event.pathway_id or "pathway"
        self._print(
            f"{Colors.BLUE}▶ PATHWAY{Colors.RESET} {Colors.BOLD}{pathway_id}{Colors.RESET}",
        )
        self._indent += 1

    def _on_execution_completed(self, event: PathwayEvent) -> None:
        self._indent = max(0, self._indent - 1)
        data = event.data or {}
        success = data.get("success", True)
        status = f"{Colors.GREEN}✓{Colors.RESET}" if success else f"{Colors.RED}✗{Colors.RESET}"
        dur = _format_duration(data.get("duration_ms"))
        self._print(
            f"{Colors.BLUE}■ PATHWAY{Colors.RESET} {status} {Colors.DIM}{dur}{Colors.RESET}"
        )

    def _on_node_started(self, event: PathwayEvent) -> None:
        self._print(f"{Colors.MAGENTA}○{Colors.RESET} {event.node_id or 'node'}")

    def _on_node_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        dur = _format_duration(data.get("duration_ms"))
        self._print(
            f"{Colors.GREEN}●{Colors.RESET} {event.node_id or 'node'} {Colors.DIM}{dur}{Colors.RESET}",
        )

    def _on_node_failed(self, event: PathwayEvent) -> None:
        self._print(
            f"{Colors.RED}●{Colors.RESET} {event.node_id or 'node'} {Colors.RED}{event.error or ''}{Colors.RESET}",
        )

    def _on_tool_called(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = data.get("tool_name", "tool")
        args = _truncate(json.dumps(data.get("args", {}), default=str), 40)
        self._print(
            f"{Colors.YELLOW}⚡{Colors.RESET} {Colors.BOLD}{tool_name}{Colors.RESET}({args})",
        )

    def _on_tool_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = data.get("tool_name", "tool")
        dur = _format_duration(data.get("duration_ms"))
        self._print(
            f"  {Colors.GREEN}✓{Colors.RESET} {tool_name} {Colors.DIM}{dur}{Colors.RESET}",
        )

    def _on_tool_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = data.get("tool_name", "tool")
        self._print(
            f"  {Colors.RED}✗{Colors.RESET} {tool_name} {Colors.RED}{event.error or ''}{Colors.RESET}",
        )


_DEBUG_DISPATCH: dict[str, Callable[[DebugHandler, PathwayEvent], None]] = {
    "execution_started": DebugHandler._on_execution_started,
    "execution_completed": DebugHandler._on_execution_completed,
    "node_started": DebugHandler._on_node_started,
    "node_completed": DebugHandler._on_node_completed,
    "node_failed": DebugHandler._on_node_failed,
    "tool_called": DebugHandler._on_tool_called,
    "tool_completed": DebugHandler._on_tool_completed,
    "tool_failed": DebugHandler._on_tool_failed,
}


# =============================================================================
# GRAPH TRACE LOGGER
# =============================================================================


class GraphTraceLogger:
    """Compact visual trace logger for pathway executions.

    Consumes ``PathwayEvent`` from the ``VMEventEmitter``.

    Example output::

        🚀 PATHWAY ai_news (exec_1234)
           │
           ├─ 🔧 search
           │     web.search({"query": "AI news"})
           │     ✓ 152ms
           │
           ├─ 🧠 summarize
           │     gpt-4o (1250→380 tok)
           │     ✓ 845ms
           │
           └─ ⚡ format ✓ 12ms

        🏁 COMPLETED 1.0s (3 nodes)
    """

    # Node type icons (compact)
    ICONS: ClassVar[dict[str, str]] = {
        "tool": "🔧",
        "llm": "🧠",
        "agent": "🤖",
        "approval": "✋",
        "for_each": "🔄",
        "pathway": "📦",
    }

    def __init__(
        self,
        use_color: bool = True,
        show_inputs: bool = False,
        json_mode: bool = False,
    ) -> None:
        self.use_color = use_color and not json_mode
        self.show_inputs = show_inputs
        self.json_mode = json_mode
        self._node_count: dict[str, int] = {}
        self._pending_nodes: dict[str, str] = {}  # exec_id -> current node_id

    def __call__(self, event: PathwayEvent) -> None:
        if self.json_mode:
            self._emit_json(event)
        else:
            handler = _TRACE_LOG_DISPATCH.get(event.event_type)
            if handler:
                handler(self, event)

    def _emit_json(self, event: PathwayEvent) -> None:
        """Emit structured JSON event for production logging."""
        data = event.data or {}
        obj: dict[str, Any] = {"event": event.event_type}
        if event.pathway_id:
            obj["pathway_id"] = event.pathway_id
        if event.node_id:
            obj["node_id"] = event.node_id
        if event.error:
            obj["error"] = event.error
        # Pull common fields from data
        for key in (
            "execution_id",
            "tool_name",
            "duration_ms",
            "model",
            "tokens_in",
            "tokens_out",
            "node_type",
            "success",
        ):
            val = data.get(key)
            if val is not None:
                obj[key] = val
        _trace_log.info(
            "event=%s %s",
            event.event_type,
            " ".join(f"{k}={v}" for k, v in obj.items() if k != "event"),
        )

    def _c(self, text: str, color: str) -> str:
        if not self.use_color:
            return text
        codes = {
            "dim": Colors.DIM,
            "bold": Colors.BOLD,
            "red": Colors.RED,
            "green": Colors.GREEN,
            "yellow": Colors.YELLOW,
            "cyan": Colors.CYAN,
            "magenta": Colors.MAGENTA,
        }
        return f"{codes.get(color, '')}{text}{Colors.RESET}"

    def _log(self, msg: str) -> None:
        _trace_log.info(msg)

    # -- Per-event handlers ---------------------------------------------------

    def _on_execution_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        exec_id = str(data.get("execution_id", "—"))
        pathway = _truncate(event.pathway_id or "pathway", 25)
        self._node_count[exec_id] = 0
        self._log("")
        self._log(f"{self._c(f'🚀 PATHWAY {pathway}', 'magenta')}{self._c(f' ({exec_id})', 'dim')}")
        self._log(self._c("   │", "dim"))

    def _on_execution_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        exec_id = str(data.get("execution_id", "—"))
        dur = _format_duration(data.get("duration_ms"))
        n = self._node_count.pop(exec_id, 0)
        success = data.get("success", True)
        if success:
            self._log(self._c(f"🏁 COMPLETED {dur} ({n} nodes)", "green"))
        else:
            err = _truncate(str(data.get("error", event.error or "")), 30)
            self._log(self._c(f"💥 FAILED {dur}: {err}", "red"))
        self._log("")

    def _on_node_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        exec_id = str(data.get("execution_id", "—"))
        node_id = event.node_id or "node"
        node_type = str(data.get("node_type", "—"))

        if exec_id in self._node_count:
            self._node_count[exec_id] += 1
        self._pending_nodes[exec_id] = node_id

        icon = self.ICONS.get(node_type, "○")
        self._log(f"{self._c('   ├─ ', 'dim')}{icon} {self._c(node_id, 'bold')}")

    def _on_node_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        dur = _format_duration(data.get("duration_ms"))
        self._log(f"{self._c('   │  ', 'dim')}{self._c(f'   ✓ {dur}', 'green')}")
        self._log(self._c("   │", "dim"))

    def _on_node_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        err = _truncate(event.error or str(data.get("error", "error")), 35)
        dur = _format_duration(data.get("duration_ms"))
        self._log(f"{self._c('   │  ', 'dim')}{self._c(f'   ✗ {dur} {err}', 'red')}")
        self._log(self._c("   │", "dim"))

    def _on_tool_called(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool = str(data.get("tool_name", "tool"))
        args = _truncate(json.dumps(data.get("args", {}), default=str), 30)
        self._log(f"{self._c(f'   │     {tool}', 'dim')}{self._c(f'({args})', 'dim')}")

    def _on_tool_completed(self, _event: PathwayEvent) -> None:
        pass  # Shown in node_completed

    def _on_tool_failed(self, event: PathwayEvent) -> None:
        err = _truncate(event.error or "error", 30)
        self._log(f"{self._c('   │     ', 'dim')}{self._c(f'✗ {err}', 'red')}")

    def _on_llm_call_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        model = str(data.get("model", "llm"))
        tokens_in = data.get("tokens_in", 0)
        tokens_out = data.get("tokens_out", 0)
        if tokens_in or tokens_out:
            self._log(self._c(f"   │     {model} ({tokens_in}→{tokens_out} tok)", "dim"))


_TRACE_LOG_DISPATCH: dict[str, Callable[[GraphTraceLogger, PathwayEvent], None]] = {
    "execution_started": GraphTraceLogger._on_execution_started,
    "execution_completed": GraphTraceLogger._on_execution_completed,
    "node_started": GraphTraceLogger._on_node_started,
    "node_completed": GraphTraceLogger._on_node_completed,
    "node_failed": GraphTraceLogger._on_node_failed,
    "tool_called": GraphTraceLogger._on_tool_called,
    "tool_completed": GraphTraceLogger._on_tool_completed,
    "tool_failed": GraphTraceLogger._on_tool_failed,
    "llm_call_completed": GraphTraceLogger._on_llm_call_completed,
}


# =============================================================================
# JSON LINE HANDLER (structured output for production)
# =============================================================================


class JsonLineHandler:
    """Write structured JSON lines to stderr for each ``PathwayEvent``.

    Use as the default handler when ``ALBUS_LOG_FORMAT=json``.  Each event
    becomes a single-line JSON object written to stderr, suitable for
    ingestion by Datadog, Loki, Splunk, or ``jq``.

    Example output::

        {"ts":"2025-01-01T00:00:00Z","event":"node_completed","node_id":"fetch",...}
    """

    def __init__(self) -> None:
        self._stream = sys.stderr

    def __call__(self, event: PathwayEvent) -> None:
        obj: dict[str, Any] = {
            "ts": datetime.now(tz=UTC).isoformat(),
            "event": event.event_type,
        }
        if event.pathway_id:
            obj["pathway_id"] = event.pathway_id
        if event.node_id:
            obj["node_id"] = event.node_id
        if event.error:
            obj["error"] = event.error
        # Flatten data into the JSON line
        data = event.data or {}
        for key in (
            "execution_id",
            "tool_name",
            "call_id",
            "duration_ms",
            "model",
            "tokens_in",
            "tokens_out",
            "node_type",
            "success",
            "error_type",
        ):
            val = data.get(key)
            if val is not None:
                obj[key] = val

        try:
            self._stream.write(json.dumps(obj, default=str) + "\n")
            self._stream.flush()
        except Exception:
            # Never break execution due to observability failures.
            pass


__all__ = [
    "Colors",
    "DebugHandler",
    "GraphTraceLogger",
    "JsonLineHandler",
]
