"""Trace logger for EB1-GNN server requests.

Captures full request/response data for analysis and debugging.
Writes JSONL files to a configurable directory with automatic rotation.

Design:
    - Async-safe for aiohttp handlers
    - Captures full input (messages, tools) and output (text, tool_calls, routing)
    - Preserves internal metadata (_selected_model, _tier, etc.) before stripping
    - Automatic daily file rotation
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class TraceLogger:
    """Async-safe trace logger for GNN server requests."""

    def __init__(
        self,
        trace_dir: str | Path | None = None,
        enabled: bool = True,
        buffer_size: int = 100,
        flush_interval_seconds: float = 5.0,
    ):
        """Initialize the trace logger.

        Args:
            trace_dir: Directory for trace files. Defaults to /var/log/eb1_traces
                or ~/eb1_traces if not writable.
            enabled: Whether tracing is enabled.
            buffer_size: Number of entries to buffer before flush.
            flush_interval_seconds: Max seconds between flushes.
        """
        self.enabled = enabled
        self.buffer_size = buffer_size
        self.flush_interval = flush_interval_seconds

        # Resolve trace directory
        if trace_dir:
            self.trace_dir = Path(trace_dir)
        else:
            # Try /var/log first, fall back to home directory
            default_dir = Path("/var/log/eb1_traces")
            if not default_dir.parent.exists() or not os.access(default_dir.parent, os.W_OK):
                default_dir = Path.home() / "eb1_traces"
            self.trace_dir = default_dir

        # State
        self._buffer: list[dict[str, Any]] = []
        self._lock = asyncio.Lock()
        self._current_file: Path | None = None
        self._current_date: str | None = None
        self._flush_task: asyncio.Task | None = None
        self._initialized = False

    async def initialize(self) -> None:
        """Initialize the trace logger (create directory, start flush task)."""
        if not self.enabled:
            return

        try:
            self.trace_dir.mkdir(parents=True, exist_ok=True)
            logger.info("Trace logging enabled: %s", self.trace_dir)
            self._initialized = True

            # Start background flush task
            self._flush_task = asyncio.create_task(self._periodic_flush())

        except Exception as e:
            logger.warning("Failed to initialize trace logging: %s", e)
            self.enabled = False

    async def shutdown(self) -> None:
        """Shutdown the trace logger (flush remaining entries)."""
        if self._flush_task:
            self._flush_task.cancel()
            try:
                await self._flush_task
            except asyncio.CancelledError:
                pass

        # Final flush
        await self._flush()

    def _get_trace_file(self) -> Path:
        """Get the current trace file path (rotates daily)."""
        today = datetime.now(timezone.utc).strftime("%Y%m%d")

        if self._current_date != today:
            self._current_date = today
            self._current_file = self.trace_dir / f"traces_{today}.jsonl"

        return self._current_file

    async def log_request(
        self,
        request_id: str,
        model: str,
        endpoint: str,
        request_body: dict[str, Any],
        response_data: dict[str, Any],
        latency_ms: float,
        error: str | None = None,
    ) -> None:
        """Log a request/response pair.

        Args:
            request_id: Unique request identifier.
            model: Requested model name.
            endpoint: API endpoint (e.g., "/v1/responses").
            request_body: Full request body.
            response_data: Full response data (including _internal fields).
            latency_ms: Request latency in milliseconds.
            error: Error message if request failed.
        """
        if not self.enabled or not self._initialized:
            return

        # Extract routing metadata (before it gets stripped)
        routing_info = {
            k: v for k, v in response_data.items()
            if k.startswith("_") and v is not None
        }

        # Extract tool calls from output
        tool_calls = []
        output = response_data.get("output", [])
        if isinstance(output, list):
            for item in output:
                if isinstance(item, dict) and item.get("type") == "function_call":
                    tool_calls.append({
                        "name": item.get("name"),
                        "arguments": item.get("arguments"),
                        "call_id": item.get("call_id"),
                    })

        # Extract text output
        output_text = response_data.get("output_text", "")
        if not output_text and isinstance(output, list):
            for item in output:
                if isinstance(item, dict) and item.get("type") == "message":
                    content = item.get("content", [])
                    if isinstance(content, list):
                        for part in content:
                            if isinstance(part, dict) and part.get("type") == "output_text":
                                output_text = part.get("text", "")
                                break

        # Build trace entry
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "request_id": request_id,
            "endpoint": endpoint,
            "model": model,
            "latency_ms": round(latency_ms, 2),
            "error": error,
            # Input
            "input": {
                "messages": self._extract_messages(request_body),
                "tools": request_body.get("tools"),
                "tool_choice": request_body.get("tool_choice"),
                "max_output_tokens": request_body.get("max_output_tokens"),
                "temperature": request_body.get("temperature"),
                "reasoning": request_body.get("reasoning"),
            },
            # Output
            "output": {
                "text": output_text[:2000] if output_text else None,  # Truncate long text
                "text_length": len(output_text) if output_text else 0,
                "tool_calls": tool_calls,
                "status": response_data.get("status"),
            },
            # Routing (eb1 internal metadata)
            "routing": routing_info,
            # Usage
            "usage": response_data.get("usage"),
        }

        async with self._lock:
            self._buffer.append(entry)

            if len(self._buffer) >= self.buffer_size:
                await self._flush_locked()

    def _extract_messages(self, body: dict[str, Any]) -> list[dict[str, Any]] | None:
        """Extract and summarize messages from request body."""
        # Try Responses API format
        input_content = body.get("input")
        if input_content:
            if isinstance(input_content, str):
                return [{"role": "user", "content_preview": input_content[:500]}]
            if isinstance(input_content, list):
                messages = []
                for item in input_content[:10]:  # Limit to first 10
                    if isinstance(item, dict):
                        msg = {"role": item.get("role", "unknown")}
                        content = item.get("content")
                        if isinstance(content, str):
                            msg["content_preview"] = content[:500]
                        elif isinstance(content, list):
                            # Extract text from content items
                            texts = []
                            for c in content[:5]:
                                if isinstance(c, dict):
                                    t = c.get("text") or c.get("input_text") or c.get("output_text")
                                    if t:
                                        texts.append(str(t)[:200])
                            msg["content_preview"] = " ".join(texts)[:500]
                        messages.append(msg)
                return messages

        # Try Chat Completions format
        messages = body.get("messages")
        if messages and isinstance(messages, list):
            return [
                {
                    "role": m.get("role", "unknown"),
                    "content_preview": str(m.get("content", ""))[:500],
                }
                for m in messages[:10]
            ]

        return None

    async def _flush(self) -> None:
        """Flush buffered entries to disk."""
        async with self._lock:
            await self._flush_locked()

    async def _flush_locked(self) -> None:
        """Flush buffered entries (must hold lock)."""
        if not self._buffer:
            return

        entries = self._buffer
        self._buffer = []

        try:
            trace_file = self._get_trace_file()
            # Run file I/O in thread pool
            await asyncio.get_event_loop().run_in_executor(
                None,
                self._write_entries,
                trace_file,
                entries,
            )
        except Exception as e:
            logger.warning("Failed to write trace entries: %s", e)
            # Put entries back in buffer (best effort)
            self._buffer = entries + self._buffer

    def _write_entries(self, trace_file: Path, entries: list[dict]) -> None:
        """Write entries to file (runs in thread pool)."""
        with open(trace_file, "a") as f:
            for entry in entries:
                f.write(json.dumps(entry, default=str) + "\n")

    async def _periodic_flush(self) -> None:
        """Periodically flush buffer."""
        while True:
            try:
                await asyncio.sleep(self.flush_interval)
                await self._flush()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.warning("Periodic flush failed: %s", e)


# Global trace logger instance
_trace_logger: TraceLogger | None = None


def get_trace_logger() -> TraceLogger:
    """Get the global trace logger instance."""
    global _trace_logger
    if _trace_logger is None:
        trace_dir = os.environ.get("EB1_TRACE_DIR")
        enabled = os.environ.get("EB1_TRACE_ENABLED", "true").lower() in ("true", "1", "yes")
        _trace_logger = TraceLogger(trace_dir=trace_dir, enabled=enabled)
    return _trace_logger


async def init_trace_logger() -> TraceLogger:
    """Initialize and return the global trace logger."""
    trace_logger = get_trace_logger()
    await trace_logger.initialize()
    return trace_logger
