"""
fw_server/tracer.py — Deep execution tracing for the fw-engine MCP server.

Records every function call, intermediate value, code path, and timing
in the fw_server pipeline.  Captures the "footprint of electrons" across:

  Layer 1: MCP tool handler calls (argument/result capture — effective
           JSON-RPC content without touching the transport layer)
  Layer 2: Dusun engine steps (7-step algorithm internals)
  Layer 3: Lens adapter dispatch (per-adapter input/output)
  Layer 4: Module-level operations (persistence, memory bank, HCP)
  Layer 5: Python logger output (all fw_server.* loggers)

Output files per session (in FW_TRACE_DIR):
  trace_<session>.jsonl  — Structured execution events (JSONL)
  logs_<session>.log     — Python logger output (human-readable)

Design constraints:
  - stdlib-only (no external dependencies)
  - Zero cost when FW_TRACE_DIR is not set (disabled)
  - KV₇ preserved: tracer is infrastructure, not analytical instrument
  - Thread-safe; uses contextvars for call-chain tracking
  - Never lets tracing errors break the server

Activation:
  Set FW_TRACE_DIR in environment or mcp.json env block:
    "env": {"FW_TRACE_DIR": ".fw_traces"}

  Each server session creates timestamped files:
    .fw_traces/trace_20260228_143000.jsonl
    .fw_traces/logs_20260228_143000.log
"""

from __future__ import annotations

import contextvars
import dataclasses
import functools
import inspect
import json
import logging
import os
import sys
import threading
import time
import traceback as tb_module
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, TextIO, TypeVar

F = TypeVar("F", bound=Callable[..., Any])

logger = logging.getLogger("fw_server.tracer")


# ── Context Variables for Call-Chain Tracking ─────────────────────────
# These track nested call IDs so every trace event knows its parent.

_call_stack: contextvars.ContextVar[List[str]] = contextvars.ContextVar(
    "fw_trace_call_stack", default=[],
)
_current_call_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "fw_trace_call_id", default=None,
)


# ── Configuration ─────────────────────────────────────────────────────

@dataclass(frozen=True)
class TraceConfig:
    """Configuration for the tracer."""

    trace_dir: Path
    max_arg_length: int = 2000       # Truncate large args in serialisation
    max_result_length: int = 15000   # Capture full dusun() v3 output (~12K)
    capture_logs: bool = True        # Route Python logs to file
    session_id: str = field(
        default_factory=lambda: datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S"),
    )


# ── Tracer Singleton ──────────────────────────────────────────────────

class Tracer:
    """Singleton deep execution tracer for fw-engine.

    All trace events are written as JSON Lines (one JSON object per line)
    to ``trace_<session>.jsonl`` inside *trace_dir*.

    The tracer is initialised once per server process via
    :func:`init_tracing`.  When uninitialised, ``Tracer.get()`` returns
    ``None`` and the ``@trace`` decorator becomes a no-op.
    """

    _instance: Optional["Tracer"] = None
    _lock = threading.Lock()

    # ── Class-level access ────────────────────────────────────────

    @classmethod
    def get(cls) -> Optional["Tracer"]:
        """Return the active tracer, or ``None`` if tracing is disabled."""
        return cls._instance

    @classmethod
    def initialize(cls, config: TraceConfig) -> "Tracer":
        """Create and install the singleton tracer."""
        with cls._lock:
            if cls._instance is not None:
                cls._instance.close()
            instance = cls(config)
            cls._instance = instance
            return instance

    @classmethod
    def shutdown(cls) -> None:
        """Shutdown, flush, and release the singleton."""
        with cls._lock:
            if cls._instance is not None:
                cls._instance.close()
                cls._instance = None

    # ── Instance ──────────────────────────────────────────────────

    def __init__(self, config: TraceConfig) -> None:
        self._config = config
        self._write_lock = threading.Lock()

        # Ensure directory exists
        config.trace_dir.mkdir(parents=True, exist_ok=True)

        # Open JSONL trace file
        self._trace_path = config.trace_dir / f"trace_{config.session_id}.jsonl"
        self._trace_file: TextIO = open(
            self._trace_path, "a", encoding="utf-8",
        )

        # Optional: Python log file handler
        self._log_handler: Optional[logging.FileHandler] = None
        if config.capture_logs:
            log_path = config.trace_dir / f"logs_{config.session_id}.log"
            self._log_handler = logging.FileHandler(
                str(log_path), encoding="utf-8",
            )
            self._log_handler.setFormatter(logging.Formatter(
                "%(asctime)s.%(msecs)03d | %(levelname)-5s | %(name)s | %(message)s",
                datefmt="%Y-%m-%d %H:%M:%S",
            ))
            # Attach to fw_server root logger so ALL sub-loggers are captured
            fw_logger = logging.getLogger("fw_server")
            fw_logger.addHandler(self._log_handler)
            fw_logger.setLevel(logging.DEBUG)

        # Session-start event
        self._emit({
            "event": "session_start",
            "session": config.session_id,
            "trace_dir": str(config.trace_dir),
            "config": {
                "max_arg_length": config.max_arg_length,
                "max_result_length": config.max_result_length,
                "capture_logs": config.capture_logs,
            },
            "python_version": sys.version,
            "pid": os.getpid(),
        })

    # ── Core emitters ─────────────────────────────────────────────

    def _emit(self, event: Dict[str, Any]) -> None:
        """Write a single trace event to the JSONL file.  Never raises."""
        event.setdefault("ts", datetime.now(timezone.utc).isoformat())
        event.setdefault("session", self._config.session_id)

        with self._write_lock:
            try:
                line = json.dumps(event, ensure_ascii=False, default=str)
                self._trace_file.write(line + "\n")
                self._trace_file.flush()
            except Exception:
                pass  # NEVER let tracing break the server

    def trace_call_start(
        self,
        func_name: str,
        module: str,
        args: Dict[str, Any],
        call_id: str,
    ) -> None:
        """Record function call entry."""
        stack = _call_stack.get([])
        # stack already contains call_id at [-1]; parent is [-2]
        parent_id = stack[-2] if len(stack) >= 2 else None
        safe_args = self._safe_serialize(args, self._config.max_arg_length)

        self._emit({
            "event": "call_start",
            "func": func_name,
            "module": module,
            "call_id": call_id,
            "parent_id": parent_id,
            "depth": max(0, len(stack) - 1),
            "args": safe_args,
        })

    def trace_call_end(
        self,
        func_name: str,
        module: str,
        call_id: str,
        duration_ms: float,
        result: Any = None,
        error: Optional[str] = None,
    ) -> None:
        """Record function call exit."""
        stack = _call_stack.get([])
        # stack still contains call_id at [-1]; parent is [-2]
        parent_id = stack[-2] if len(stack) >= 2 else None

        event: Dict[str, Any] = {
            "event": "call_end",
            "func": func_name,
            "module": module,
            "call_id": call_id,
            "parent_id": parent_id,
            "depth": max(0, len(stack) - 1),
            "duration_ms": round(duration_ms, 3),
        }

        if error is not None:
            event["error"] = error[:2000]
        else:
            event["result_type"] = type(result).__name__
            event["result_summary"] = self._summarize(
                result, self._config.max_result_length,
            )

        self._emit(event)

    def trace_intermediate(
        self,
        func_name: str,
        label: str,
        value: Any,
    ) -> None:
        """Record an intermediate computed value within a function."""
        self._emit({
            "event": "intermediate",
            "func": func_name,
            "label": label,
            "value": self._safe_serialize(
                value, self._config.max_result_length,
            ),
            "call_id": _current_call_id.get(None),
        })

    def trace_event(self, event_type: str, **data: Any) -> None:
        """Record a custom / ad-hoc event."""
        self._emit({
            "event": event_type,
            "call_id": _current_call_id.get(None),
            **{k: self._safe_serialize(v, 1000) for k, v in data.items()},
        })

    # ── Serialisation helpers ─────────────────────────────────────

    def _summarize(self, value: Any, max_len: int) -> Any:
        """Create a truncated summary of a return value."""
        if value is None:
            return None
        if isinstance(value, (int, float, bool)):
            return value
        if isinstance(value, str):
            if len(value) > max_len:
                return value[:max_len] + f"...[truncated, total={len(value)}]"
            return value
        if isinstance(value, dict):
            out: Dict[str, Any] = {}
            for k, v in list(value.items())[:30]:
                out[str(k)] = self._summarize(v, max_len)
            if len(value) > 30:
                out["__truncated__"] = f"{len(value)} total keys"
            return out
        if isinstance(value, (list, tuple)):
            if len(value) > 20:
                return {
                    "__type__": type(value).__name__,
                    "__len__": len(value),
                    "first_5": [self._summarize(x, max_len) for x in value[:5]],
                }
            return [self._summarize(x, max_len) for x in value]
        # FIX-04: dataclass-aware serialization (LensResult, QualityReport, etc.)
        if dataclasses.is_dataclass(value) and not isinstance(value, type):
            return self._summarize(dataclasses.asdict(value), max_len)
        if hasattr(value, 'to_dict') and callable(value.to_dict):
            return self._summarize(value.to_dict(), max_len)
        r = repr(value)
        return r[:max_len] + "...[truncated]" if len(r) > max_len else r

    def _safe_serialize(self, value: Any, max_len: int) -> Any:
        """Safely convert arbitrary values for JSON output."""
        if value is None:
            return None
        if isinstance(value, (int, float, bool)):
            return value
        if isinstance(value, str):
            if len(value) > max_len:
                return value[:max_len] + f"...[truncated, total={len(value)}]"
            return value
        if isinstance(value, bytes):
            return f"<bytes len={len(value)}>"
        if isinstance(value, Path):
            return str(value)
        if isinstance(value, dict):
            return {
                str(k): self._safe_serialize(v, max_len)
                for k, v in list(value.items())[:50]
            }
        if isinstance(value, (list, tuple)):
            if len(value) > 50:
                return {
                    "__type__": type(value).__name__,
                    "__len__": len(value),
                    "sample": [
                        self._safe_serialize(x, max_len) for x in value[:10]
                    ],
                }
            return [self._safe_serialize(x, max_len) for x in value]
        # FIX-04: dataclass-aware serialization (LensResult, QualityReport, etc.)
        if dataclasses.is_dataclass(value) and not isinstance(value, type):
            try:
                return self._safe_serialize(dataclasses.asdict(value), max_len)
            except Exception:
                pass
        if hasattr(value, 'to_dict') and callable(value.to_dict):
            try:
                return self._safe_serialize(value.to_dict(), max_len)
            except Exception:
                pass
        try:
            r = repr(value)
            return r[:max_len] if len(r) > max_len else r
        except Exception:
            return f"<{type(value).__name__}>"

    # ── Lifecycle ─────────────────────────────────────────────────

    def close(self) -> None:
        """Flush and close all trace files."""
        self._emit({"event": "session_end"})

        with self._write_lock:
            try:
                self._trace_file.close()
            except Exception:
                pass

        if self._log_handler:
            fw_logger = logging.getLogger("fw_server")
            fw_logger.removeHandler(self._log_handler)
            try:
                self._log_handler.close()
            except Exception:
                pass


# ══════════════════════════════════════════════════════════════════════
# Public API: @trace decorator
# ══════════════════════════════════════════════════════════════════════

def trace(func: F) -> F:
    """Decorator that records function entry, exit, args, result, duration.

    **Zero-cost** when tracing is disabled (``Tracer.get()`` returns
    ``None``).  Safe to leave permanently applied in production code.
    Supports both sync and async functions.

    Usage::

        @trace
        def my_function(x, y):
            return x + y

        @trace
        async def my_async_function(x, y):
            return x + y
    """
    func_name = func.__qualname__
    module = func.__module__ or ""

    def _setup(args: Any, kwargs: Any):
        """Common pre-call setup; returns (tracer, call_id, tok_stack, tok_id, t0) or None."""
        tracer = Tracer.get()
        if tracer is None:
            return None
        call_id = uuid.uuid4().hex[:12]
        arg_dict = _build_arg_dict(func, args, kwargs)
        stack = _call_stack.get([])
        new_stack = stack + [call_id]
        tok_stack = _call_stack.set(new_stack)
        tok_id = _current_call_id.set(call_id)
        tracer.trace_call_start(func_name, module, arg_dict, call_id)
        t0 = time.perf_counter()
        return tracer, call_id, tok_stack, tok_id, t0

    if inspect.iscoroutinefunction(func):
        @functools.wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            ctx = _setup(args, kwargs)
            if ctx is None:
                return await func(*args, **kwargs)
            tracer, call_id, tok_stack, tok_id, t0 = ctx
            try:
                result = await func(*args, **kwargs)
                duration_ms = (time.perf_counter() - t0) * 1000
                tracer.trace_call_end(
                    func_name, module, call_id, duration_ms, result=result,
                )
                return result
            except Exception as exc:
                duration_ms = (time.perf_counter() - t0) * 1000
                tracer.trace_call_end(
                    func_name, module, call_id, duration_ms,
                    error=f"{type(exc).__name__}: {exc}\n{tb_module.format_exc()}",
                )
                raise
            finally:
                _call_stack.reset(tok_stack)
                _current_call_id.reset(tok_id)
        return async_wrapper  # type: ignore[return-value]

    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        ctx = _setup(args, kwargs)
        if ctx is None:
            return func(*args, **kwargs)
        tracer, call_id, tok_stack, tok_id, t0 = ctx
        try:
            result = func(*args, **kwargs)
            duration_ms = (time.perf_counter() - t0) * 1000
            tracer.trace_call_end(
                func_name, module, call_id, duration_ms, result=result,
            )
            return result
        except Exception as exc:
            duration_ms = (time.perf_counter() - t0) * 1000
            tracer.trace_call_end(
                func_name, module, call_id, duration_ms,
                error=f"{type(exc).__name__}: {exc}\n{tb_module.format_exc()}",
            )
            raise
        finally:
            _call_stack.reset(tok_stack)
            _current_call_id.reset(tok_id)

    return wrapper  # type: ignore[return-value]


# ══════════════════════════════════════════════════════════════════════
# Public API: record() — inline intermediate value capture
# ══════════════════════════════════════════════════════════════════════

def record(label: str, value: Any) -> Any:
    """Record an intermediate value.  Returns *value* unchanged for inline use.

    When tracing is disabled this is a pure identity function (zero cost).

    Usage::

        score = record("composite_score", bileshke.compute(results))
    """
    tracer = Tracer.get()
    if tracer is not None:
        frame = inspect.currentframe()
        caller = frame.f_back if frame else None
        if caller and hasattr(caller.f_code, "co_qualname"):
            func_name = caller.f_code.co_qualname
        elif caller:
            func_name = caller.f_code.co_name
        else:
            func_name = "unknown"
        tracer.trace_intermediate(func_name, label, value)
    return value


# ══════════════════════════════════════════════════════════════════════
# Public API: init / shutdown
# ══════════════════════════════════════════════════════════════════════

def init_tracing(trace_dir: Optional[str] = None) -> Optional[Tracer]:
    """Initialise tracing from an explicit path or ``FW_TRACE_DIR`` env var.

    Returns the :class:`Tracer` instance if activated, else ``None``.
    Safe to call multiple times — reinitialises cleanly.
    """
    trace_path = trace_dir or os.environ.get("FW_TRACE_DIR")
    if not trace_path:
        return None

    config = TraceConfig(trace_dir=Path(trace_path))
    tracer = Tracer.initialize(config)

    logger.info(
        "Tracing initialised — session=%s  dir=%s",
        config.session_id,
        trace_path,
    )
    return tracer


def shutdown_tracing() -> None:
    """Shutdown the tracer, flush all files, and release resources."""
    Tracer.shutdown()


# ── Internal helpers ──────────────────────────────────────────────────

def _build_arg_dict(func: Callable, args: tuple, kwargs: dict) -> dict:
    """Map positional + keyword args to parameter names."""
    try:
        sig = inspect.signature(func)
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        result = dict(bound.arguments)
        result.pop("self", None)
        result.pop("cls", None)
        return result
    except (ValueError, TypeError):
        return {"args": list(args), "kwargs": kwargs}
