#!/usr/bin/env python3
"""
trace_collector.py - SER Routing Trace Event Collector

Implements the trace event format from docs/SER_TRACE_SPEC_V1.md.

Provides a context-manager based API for emitting structured trace events
during Se-based routing operations.

Usage:
    from flowtorch.trace_collector import TraceCollector

    collector = TraceCollector(output_dir="receipts/ser_traces")

    with collector.trace("unified_se_router", input_info={"label": "test"}) as t:
        t.emit_se_computed(se=0.187, components={"H_norm": 0.65, "C": 1.0, "D": 0.29})
        t.emit_threshold_check("Se < SE_LOW", 0.187, 0.20, True)
        t.emit_routing_decision(
            backend="ibm",
            caps={"qaoa_layers": 2},
            se=0.187,
            regime="low_se_structured",
            rationale="Low Se → QAOA"
        )
"""

from __future__ import annotations

import json
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional


# Schema version
TRACE_EVENT_SCHEMA = "helix:ser_trace_event:v1"
TRACE_SUMMARY_SCHEMA = "helix:ser_trace_summary:v1"

# Standard thresholds
SE_LOW = 0.20
SE_HIGH = 0.70


class TraceContext:
    """
    Active trace context for emitting events.

    Created by TraceCollector.trace() context manager.
    """

    def __init__(
        self,
        correlation_id: str,
        source: str,
        output_file: Path,
        input_info: Optional[Dict] = None,
    ):
        self.correlation_id = correlation_id
        self.source = source
        self.output_file = output_file
        self.input_info = input_info or {}
        self.sequence = 0
        self.events: List[Dict] = []
        self.start_ts: Optional[str] = None
        self._closed = False

        # Tracking for summary
        self.se: Optional[float] = None
        self.regime: Optional[str] = None
        self.backend: Optional[str] = None

    def _emit(self, event_type: str, **fields) -> Dict:
        """Emit a trace event."""
        if self._closed:
            raise RuntimeError("Trace already closed")

        ts = datetime.now(timezone.utc).isoformat()

        event = {
            "schema": TRACE_EVENT_SCHEMA,
            "event": event_type,
            "ts": ts,
            "correlation_id": self.correlation_id,
            "sequence": self.sequence,
            **fields,
        }

        self.sequence += 1
        self.events.append(event)

        # Write to file immediately (append mode)
        with open(self.output_file, "a") as f:
            f.write(json.dumps(event) + "\n")

        return event

    def emit_trace_start(self) -> Dict:
        """Emit trace_start event."""
        self.start_ts = datetime.now(timezone.utc).isoformat()
        return self._emit(
            "trace_start",
            source=self.source,
            input=self.input_info,
        )

    def emit_se_computed(
        self,
        se: float,
        components: Dict[str, float],
        formula: str = "Se = H_norm × C × D",
    ) -> Dict:
        """Emit se_computed event."""
        self.se = se
        return self._emit(
            "se_computed",
            se=round(se, 6),
            components={k: round(v, 4) for k, v in components.items()},
            formula=formula,
        )

    def emit_threshold_check(
        self,
        check: str,
        lhs: float,
        rhs: float,
        result: bool,
    ) -> Dict:
        """Emit threshold_check event."""
        return self._emit(
            "threshold_check",
            check=check,
            lhs=round(lhs, 4),
            rhs=round(rhs, 4),
            result=result,
        )

    def emit_regime_classified(
        self,
        regime: str,
        se: float,
        thresholds: Optional[Dict[str, float]] = None,
    ) -> Dict:
        """Emit regime_classified event."""
        self.regime = regime
        return self._emit(
            "regime_classified",
            regime=regime,
            se=round(se, 6),
            thresholds=thresholds or {"SE_LOW": SE_LOW, "SE_HIGH": SE_HIGH},
        )

    def emit_routing_decision(
        self,
        backend: str,
        caps: Dict[str, Any],
        se: float,
        regime: str,
        rationale: str,
        block_id: Optional[str] = None,
        decision_path: Optional[List[str]] = None,
        thresholds: Optional[Dict[str, float]] = None,
    ) -> Dict:
        """Emit routing_decision event."""
        self.backend = backend
        self.regime = regime
        self.se = se

        event = {
            "backend": backend,
            "caps": caps,
            "se": round(se, 6),
            "regime": regime,
            "thresholds": thresholds or {"SE_LOW": SE_LOW, "SE_HIGH": SE_HIGH},
            "rationale": rationale,
        }

        if block_id:
            event["block_id"] = block_id

        if decision_path:
            event["decision_path"] = decision_path

        return self._emit("routing_decision", **event)

    def emit_solve_start(
        self,
        backend: str,
        caps: Dict[str, Any],
    ) -> Dict:
        """Emit solve_start event."""
        return self._emit(
            "solve_start",
            backend=backend,
            caps=caps,
        )

    def emit_solve_complete(
        self,
        backend: str,
        runtime_ms: float,
        energy: Optional[float] = None,
        bitstring_hash: Optional[str] = None,
    ) -> Dict:
        """Emit solve_complete event."""
        event = {
            "backend": backend,
            "runtime_ms": round(runtime_ms, 3),
        }
        if energy is not None:
            event["energy"] = round(energy, 4)
        if bitstring_hash:
            event["bitstring_hash"] = bitstring_hash

        return self._emit("solve_complete", **event)

    def emit_error(
        self,
        error_type: str,
        message: str,
        fallback_backend: Optional[str] = None,
    ) -> Dict:
        """Emit error event."""
        event = {
            "error_type": error_type,
            "message": message,
        }
        if fallback_backend:
            event["fallback_backend"] = fallback_backend

        return self._emit("error", **event)

    def emit_trace_end(
        self,
        success: bool,
        total_runtime_ms: Optional[float] = None,
    ) -> Dict:
        """Emit trace_end event and close the trace."""
        event = self._emit(
            "trace_end",
            total_runtime_ms=round(total_runtime_ms, 3) if total_runtime_ms else None,
            backend_used=self.backend,
            se=round(self.se, 6) if self.se else None,
            success=success,
        )
        self._closed = True
        return event

    def get_summary(self) -> Dict:
        """Generate trace summary."""
        return {
            "schema": TRACE_SUMMARY_SCHEMA,
            "correlation_id": self.correlation_id,
            "timestamp": self.start_ts,
            "source": self.source,
            "se": round(self.se, 6) if self.se else None,
            "regime": self.regime,
            "backend": self.backend,
            "event_count": len(self.events),
            "trace_file": self.output_file.name,
        }


class TraceCollector:
    """
    SER Routing Trace Collector.

    Manages trace lifecycle and file output.
    """

    def __init__(
        self,
        output_dir: Optional[str] = None,
        enabled: bool = True,
    ):
        """
        Initialize trace collector.

        Args:
            output_dir: Directory for trace files (default: receipts/ser_traces)
            enabled: Whether tracing is enabled (allows global disable)
        """
        self.enabled = enabled
        if output_dir:
            self.output_dir = Path(output_dir)
        else:
            self.output_dir = Path(__file__).parent.parent / "receipts" / "ser_traces"

        # Ensure output directory exists
        if self.enabled:
            self.output_dir.mkdir(parents=True, exist_ok=True)

    @contextmanager
    def trace(
        self,
        source: str,
        input_info: Optional[Dict] = None,
        correlation_id: Optional[str] = None,
    ):
        """
        Context manager for a trace session.

        Usage:
            with collector.trace("my_router", {"label": "test"}) as t:
                t.emit_se_computed(...)
                t.emit_routing_decision(...)

        Args:
            source: Name of the component emitting the trace
            input_info: Optional dict describing the input
            correlation_id: Optional correlation ID (auto-generated if not provided)

        Yields:
            TraceContext for emitting events
        """
        if not self.enabled:
            # Return a no-op context when disabled
            yield _NoOpTraceContext()
            return

        cid = correlation_id or uuid.uuid4().hex[:16]
        output_file = self.output_dir / f"trace_{cid}.jsonl"

        ctx = TraceContext(
            correlation_id=cid,
            source=source,
            output_file=output_file,
            input_info=input_info,
        )

        import time
        start_time = time.perf_counter()

        try:
            ctx.emit_trace_start()
            yield ctx
            runtime_ms = (time.perf_counter() - start_time) * 1000
            ctx.emit_trace_end(success=True, total_runtime_ms=runtime_ms)

        except Exception as e:
            runtime_ms = (time.perf_counter() - start_time) * 1000
            ctx.emit_error(
                error_type=type(e).__name__,
                message=str(e),
            )
            ctx.emit_trace_end(success=False, total_runtime_ms=runtime_ms)
            raise

        finally:
            # Write summary file
            summary_file = self.output_dir / f"trace_{cid}_summary.json"
            summary = ctx.get_summary()
            summary["total_runtime_ms"] = runtime_ms
            summary["success"] = not ctx._closed or ctx.events[-1].get("success", False)
            summary_file.write_text(json.dumps(summary, indent=2))

    def start_trace(
        self,
        source: str,
        input_info: Optional[Dict] = None,
        correlation_id: Optional[str] = None,
    ) -> TraceContext:
        """
        Start a trace manually (non-context-manager style).

        Remember to call ctx.emit_trace_end() when done!

        Args:
            source: Name of the component emitting the trace
            input_info: Optional dict describing the input
            correlation_id: Optional correlation ID

        Returns:
            TraceContext for emitting events
        """
        if not self.enabled:
            return _NoOpTraceContext()

        cid = correlation_id or uuid.uuid4().hex[:16]
        output_file = self.output_dir / f"trace_{cid}.jsonl"

        ctx = TraceContext(
            correlation_id=cid,
            source=source,
            output_file=output_file,
            input_info=input_info,
        )

        ctx.emit_trace_start()
        return ctx


class _NoOpTraceContext:
    """No-op trace context when tracing is disabled."""

    def emit_trace_start(self): pass
    def emit_se_computed(self, **kw): pass
    def emit_threshold_check(self, *a, **kw): pass
    def emit_regime_classified(self, *a, **kw): pass
    def emit_routing_decision(self, **kw): pass
    def emit_solve_start(self, *a, **kw): pass
    def emit_solve_complete(self, *a, **kw): pass
    def emit_error(self, *a, **kw): pass
    def emit_trace_end(self, *a, **kw): pass
    def get_summary(self): return {}


# Module-level default collector instance
_default_collector: Optional[TraceCollector] = None


def get_collector(output_dir: Optional[str] = None) -> TraceCollector:
    """Get or create the default trace collector."""
    global _default_collector
    if _default_collector is None:
        _default_collector = TraceCollector(output_dir=output_dir)
    return _default_collector


def enable_tracing(enabled: bool = True, output_dir: Optional[str] = None):
    """Enable or disable tracing globally."""
    collector = get_collector(output_dir)
    collector.enabled = enabled


def is_tracing_enabled() -> bool:
    """Check if tracing is currently enabled."""
    collector = get_collector()
    return collector.enabled
