import sys
from typing import Any

from agenx.trace.spans.models import (
    StepAttributes,
    LLMAttributes,
    Span,
    RetrieverAttributes,
    ToolAttributes,
    GuardAttributes,
    AgentAttributes,
    SpanStatus
)
from agenx._internal import duration_ms

_GREEN = "\033[92m"
_RED = "\033[91m"
_CYAN = "\033[96m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
_RESET = "\033[0m"


def _c(text: str, code: str) -> str:
    """Apply colour only when stdout is a real terminal, not a pipe or file."""
    if sys.stdout.isatty():
        return f"{code}{text}{_RESET}"
    return text


def _fmt_duration(ms: float | None) -> str:
    if ms is None:
        return "?"
    if ms < 1:
        return f"{ms * 1000:.0f}µs"
    if ms < 1000:
        return f"{ms:.1f}ms"
    return f"{ms / 1000:.2f}s"


def _notable_attrs(span: Span) -> list[tuple[str, Any]]:
    """Return only the attribute fields worth printing."""
    attrs = span.attributes
    pairs = []

    if isinstance(attrs, StepAttributes):
        if attrs.input is not None:
            pairs.append(("input",  attrs.input))
        if attrs.output is not None:
            pairs.append(("output", attrs.output))

    elif isinstance(attrs, LLMAttributes):
        if attrs.model:
            pairs.append(("model",    attrs.model))
        if attrs.provider:
            pairs.append(("provider", attrs.provider))
        if attrs.tokens_total:
            pairs.append(("tokens",   attrs.tokens_total))
        if attrs.output:
            preview = attrs.output[:80] + \
                "…" if len(attrs.output) > 80 else attrs.output
            pairs.append(("output", preview))

    elif isinstance(attrs, RetrieverAttributes):
        if attrs.query:
            pairs.append(("query",    attrs.query[:60]))
        if attrs.returned:
            pairs.append(("returned", attrs.returned))
        if attrs.top_k:
            pairs.append(("top_k",    attrs.top_k))

    elif isinstance(attrs, ToolAttributes):
        if attrs.tool_name:
            pairs.append(("tool",   attrs.tool_name))
        if attrs.input is not None:
            pairs.append(("input",  attrs.input))
        if attrs.output is not None:
            pairs.append(("output", attrs.output))

    elif isinstance(attrs, GuardAttributes):
        if attrs.guard_name:
            pairs.append(("guard",  attrs.guard_name))
        pairs.append(("passed", attrs.passed))
        if not attrs.passed and attrs.reason:
            pairs.append(("reason", attrs.reason))

    elif isinstance(attrs, AgentAttributes):
        if attrs.iterations:
            pairs.append(("iterations", attrs.iterations))
        if attrs.tools:
            pairs.append(("tools",      attrs.tools))
        if attrs.input is not None:
            pairs.append(("input",  attrs.input))
        if attrs.output is not None:
            pairs.append(("output", attrs.output))

    if hasattr(attrs, "metadata"):
        for k, v in attrs.metadata.items():
            pairs.append((k, v))

    return pairs


class StdoutExporter:
    def __init__(
        self,
        indent: bool = True,
        show_ids: bool = True,
        min_duration_ms: float = 0.0,
    ) -> None:
        self.indent = indent
        self.show_ids = show_ids
        self.min_duration_ms = min_duration_ms

    def __call__(self, span: Span) -> None:
        try:
            self._print(span)
        except Exception:
            pass

    def _print(self, span: Span) -> None:
        dur_ms = duration_ms(
            span.start_time, span.end_time) if span.end_time else None

        if dur_ms is not None and dur_ms < self.min_duration_ms:
            return

        indent = "  " if (self.indent and span.parent_id) else ""

        icon = _c("✓", _GREEN) if span.status == SpanStatus.OK else _c("✗", _RED)
        type_tag = _c(f"[{span.type.value.upper()}]", _CYAN)
        name = _c(span.name, _BOLD)
        dur = _c(_fmt_duration(dur_ms), _DIM)

        header = f"{indent}{icon} {type_tag} {name}  {dur}"

        if self.show_ids:
            header += f"  {_c('trace=' + span.trace_id[:8], _DIM)}"
            header += f"  {_c('span='  + span.span_id[:8],  _DIM)}"

        print(header)

        pairs = _notable_attrs(span)
        if pairs:
            attr_str = "  ".join(f"{k}={v!r}" for k, v in pairs)
            print(f"{indent}  {_c(attr_str, _DIM)}")

        if span.error:
            print(f"{indent}  {_c('error: ' + repr(span.error), _RED)}")

        for event in span.events:
            import datetime
            dt = datetime.datetime.fromtimestamp(
                event.timestamp.wall_ns / 1_000_000_000,
                tz=datetime.timezone.utc
            )
            time_str = dt.strftime("%H:%M:%S.") + \
                f"{dt.microsecond // 1000:03d}"
            ev_attrs = ""
            if event.attributes:
                ev_attrs = "  " + "  ".join(
                    f"{k}={v!r}" for k, v in event.attributes.items()
                )
            print(f"{indent}  {_c('└─', _DIM)} {event.name} @ {time_str}{ev_attrs}")
