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

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

from core.observability.events import (
    Event,
    LLMRequestEvent,
    LLMResponseEvent,
    NodeCompletedEvent,
    NodeFailedEvent,
    NodeStartedEvent,
    PathwayCompletedEvent,
    PathwayCreatedEvent,
    PathwayStartedEvent,
    ToolCalledEvent,
    ToolCompletedEvent,
    ToolFailedEvent,
    TurnCompletedEvent,
    TurnFailedEvent,
    TurnStartedEvent,
)

_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]}..."


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


class DebugHandler:
    """Pretty-prints events to stdout for debugging."""

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

    def __call__(self, event: Event) -> None:
        method = getattr(self, f"_handle_{event.type.value}", None)
        if method:
            method(event)
        elif self.verbose:
            self._print_generic(event)

    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, *args: object, **kwargs: object) -> None:
        pass

    def _print_generic(self, event: Event) -> None:
        self._print(f"{Colors.DIM}[{event.type.value}]{Colors.RESET}")

    def _handle_turn_started(self, event: TurnStartedEvent) -> None:
        self._print(f"{Colors.GREEN}┌─ TURN STARTED{Colors.RESET}")
        self._print(f"{Colors.GREEN}│{Colors.RESET} User: {_truncate(event.message, 60)}")
        self._indent += 1

    def _handle_turn_completed(self, event: TurnCompletedEvent) -> None:
        self._indent = max(0, self._indent - 1)
        self._print(
            f"{Colors.GREEN}└─ TURN COMPLETED{Colors.RESET} {Colors.DIM}{event.duration_ms:.0f}ms{Colors.RESET}",
        )

    def _handle_turn_failed(self, event: TurnFailedEvent) -> None:
        self._indent = max(0, self._indent - 1)
        self._print(
            f"{Colors.RED}└─ TURN FAILED{Colors.RESET} {Colors.RED}{event.error}{Colors.RESET}",
        )

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

    def _handle_pathway_completed(self, event: PathwayCompletedEvent) -> None:
        self._indent = max(0, self._indent - 1)
        status = (
            f"{Colors.GREEN}✓{Colors.RESET}" if event.success else f"{Colors.RED}✗{Colors.RESET}"
        )
        self._print(f"{Colors.BLUE}■ PATHWAY{Colors.RESET} {status}")

    def _handle_node_started(self, event: NodeStartedEvent) -> None:
        self._print(f"{Colors.MAGENTA}○{Colors.RESET} {event.node_id}")

    def _handle_node_completed(self, event: NodeCompletedEvent) -> None:
        self._print(
            f"{Colors.GREEN}●{Colors.RESET} {event.node_id} {Colors.DIM}{event.duration_ms:.0f}ms{Colors.RESET}",
        )

    def _handle_node_failed(self, event: NodeFailedEvent) -> None:
        self._print(
            f"{Colors.RED}●{Colors.RESET} {event.node_id} {Colors.RED}{event.error}{Colors.RESET}",
        )

    def _handle_tool_called(self, event: ToolCalledEvent) -> None:
        args = _truncate(json.dumps(event.tool_args), 40) if event.tool_args else ""
        self._print(
            f"{Colors.YELLOW}⚡{Colors.RESET} {Colors.BOLD}{event.tool_name}{Colors.RESET}({args})",
        )

    def _handle_tool_completed(self, event: ToolCompletedEvent) -> None:
        self._print(
            f"  {Colors.GREEN}✓{Colors.RESET} {event.tool_name} {Colors.DIM}{event.duration_ms:.0f}ms{Colors.RESET}",
        )

    def _handle_tool_failed(self, event: ToolFailedEvent) -> None:
        self._print(
            f"  {Colors.RED}✗{Colors.RESET} {event.tool_name} {Colors.RED}{event.error}{Colors.RESET}",
        )


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


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"


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

    Shows node-by-node execution with minimal width for small terminals.

    Example output:
        🤖 TURN host ───────────────────────
           "Build a pipeline that..."

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

        🏁 COMPLETED 1.0s (3 nodes)
        ✅ TURN 1.2s
    """

    # Node type icons (compact)
    ICONS: ClassVar[dict[str, str]] = {
        "tool": "🔧",
        "llm": "🧠",
        "agent": "🤖",
        "transform": "⚡",
        "gate": "🚦",
        "loop": "🔄",
    }

    def __init__(self, use_color: bool = True, show_inputs: bool = False, json_mode: bool = False):
        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: Event) -> None:
        if self.json_mode:
            self._emit_json(event)
        else:
            method = getattr(self, f"_trace_{event.type.value}", None)
            if method:
                method(event)

    def _emit_json(self, event: Event) -> None:
        """Emit structured JSON event for production logging."""
        data: dict[str, Any] = {"event": event.type.value}
        for attr in ("agent_id", "thread_id", "execution_id", "pathway_id", "node_id"):
            val = getattr(event, attr, None)
            if val:
                data[attr] = val
        for attr in (
            "message",
            "error",
            "duration_ms",
            "success",
            "tool_name",
            "model",
            "tokens_in",
            "tokens_out",
            "node_type",
        ):
            val = getattr(event, attr, None)
            if val is not None:
                data[attr] = val
        _trace_log.info(
            "event=%s %s",
            event.type.value,
            " ".join(f"{k}={v}" for k, v in data.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)

    # Turn events
    def _trace_turn_started(self, event: TurnStartedEvent) -> None:
        agent = event.agent_id or "agent"
        msg = _truncate(event.message, 50) if event.message else ""
        self._log("")
        self._log(f"{self._c(f'🤖 TURN {agent} ', 'cyan')}{self._c('─' * 30, 'dim')}")
        if msg:
            self._log(self._c(f'   "{msg}"', "dim"))

    def _trace_turn_completed(self, event: TurnCompletedEvent) -> None:
        dur = _format_duration(event.duration_ms)
        self._log(self._c(f"✅ TURN {dur}", "green"))
        self._log("")

    def _trace_turn_failed(self, event: TurnFailedEvent) -> None:
        err = _truncate(event.error or "unknown", 40)
        self._log(self._c(f"❌ TURN FAILED: {err}", "red"))
        self._log("")

    # Pathway events
    def _trace_pathway_created(self, event: PathwayCreatedEvent) -> None:
        name = _truncate(event.pathway_name or event.pathway_id or "pathway", 30)
        self._log(f"   ✨ Created {self._c(name, 'bold')} ({event.node_count} nodes)")

    def _trace_pathway_started(self, event: PathwayStartedEvent) -> None:
        exec_id = event.execution_id or "—"
        pathway = _truncate(event.pathway_name or 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 _trace_pathway_completed(self, event: PathwayCompletedEvent) -> None:
        dur = _format_duration(event.duration_ms)
        exec_id = event.execution_id or "—"
        n = self._node_count.pop(exec_id, 0)
        if event.success:
            self._log(self._c(f"🏁 COMPLETED {dur} ({n} nodes)", "green"))
        else:
            err = _truncate(getattr(event, "error", "") or "", 30)
            self._log(self._c(f"💥 FAILED {dur}: {err}", "red"))

    # Node events
    def _trace_node_started(self, event: NodeStartedEvent) -> None:
        exec_id = event.execution_id or "—"
        node_id = event.node_id or "node"
        node_type = event.node_type or "—"

        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 _trace_node_completed(self, event: NodeCompletedEvent) -> None:
        dur = _format_duration(event.duration_ms)
        self._log(f"{self._c('   │  ', 'dim')}{self._c(f'   ✓ {dur}', 'green')}")
        self._log(self._c("   │", "dim"))

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

    # Tool events
    def _trace_tool_called(self, event: ToolCalledEvent) -> None:
        tool = event.tool_name or "tool"
        args = _truncate(json.dumps(event.tool_args or {}), 30)
        self._log(f"{self._c(f'   │     {tool}', 'dim')}{self._c(f'({args})', 'dim')}")

    def _trace_tool_completed(self, event: ToolCompletedEvent) -> None:
        pass  # Shown in node_completed

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

    # LLM events
    def _trace_llm_request(self, event: LLMRequestEvent) -> None:
        model = event.model or "llm"
        self._log(self._c(f"   │     {model}", "dim"))

    def _trace_llm_response(self, event: LLMResponseEvent) -> None:
        tokens_in = event.tokens_in or 0
        tokens_out = event.tokens_out or 0
        if tokens_in or tokens_out:
            self._log(self._c(f"   │     ({tokens_in}→{tokens_out} tok)", "dim"))


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


class JsonLineHandler:
    """Write structured JSON lines to stderr for each observability event.

    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:
        import sys

        self._stream = sys.stderr

    def __call__(self, event: Event) -> None:
        data: dict[str, Any] = {
            "ts": datetime.now(tz=UTC).isoformat(),
            "event": event.type.value,
        }
        # Attach standard context fields.
        for attr in (
            "agent_id",
            "thread_id",
            "execution_id",
            "pathway_id",
            "node_id",
        ):
            val = getattr(event, attr, None)
            if val:
                data[attr] = val
        # Attach event-specific fields.
        for attr in (
            "message",
            "error",
            "duration_ms",
            "success",
            "tool_name",
            "tool_args",
            "model",
            "tokens_in",
            "tokens_out",
            "node_type",
            "pathway_name",
            "node_count",
        ):
            val = getattr(event, attr, None)
            if val is not None:
                data[attr] = val

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


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