"""
Lightweight observability layer.

Tracks per-request metrics and persists them to logs/logs.json.
Metrics collected:
  - timestamp
  - model_used
  - backend
  - prompt_length (chars)
  - response_length (chars)
  - estimated_prompt_tokens
  - estimated_response_tokens
  - latency_ms
  - was_refused (safety refusal)
  - violation_type (if refused)
  - session_id
"""

from __future__ import annotations

import json
import logging
import os
import time
import uuid
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import List, Optional

logger = logging.getLogger(__name__)

LOGS_PATH = Path("logs/logs.json")


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------

@dataclass
class RequestLog:
    session_id: str
    model_used: str
    backend: str
    prompt_length: int
    response_length: int
    estimated_prompt_tokens: int
    estimated_response_tokens: int
    latency_ms: float
    was_refused: bool
    violation_type: Optional[str]
    timestamp: str = field(default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
    request_id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])

    def to_dict(self) -> dict:
        return asdict(self)


# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------

def estimate_tokens(text: str) -> int:
    """
    Rough token estimate: ~4 characters per token (standard approximation).
    Good enough for observability; not a replacement for a proper tokenizer.
    """
    return max(1, len(text) // 4)


# ---------------------------------------------------------------------------
# Observability manager
# ---------------------------------------------------------------------------

class ObservabilityManager:
    """
    Collects and persists request logs.
    Thread-safe for single-process Streamlit apps.
    """

    def __init__(self, log_path: Path = LOGS_PATH):
        self.log_path = log_path
        self._logs: List[RequestLog] = []
        self._session_id: str = str(uuid.uuid4())[:8]
        self._ensure_log_dir()
        self._load_existing_logs()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def log_request(
        self,
        model_used: str,
        backend: str,
        prompt: str,
        response: str,
        latency_ms: float,
        was_refused: bool = False,
        violation_type: Optional[str] = None,
    ) -> RequestLog:
        """Record a single request/response pair."""
        log = RequestLog(
            session_id=self._session_id,
            model_used=model_used,
            backend=backend,
            prompt_length=len(prompt),
            response_length=len(response),
            estimated_prompt_tokens=estimate_tokens(prompt),
            estimated_response_tokens=estimate_tokens(response),
            latency_ms=round(latency_ms, 2),
            was_refused=was_refused,
            violation_type=violation_type,
        )
        self._logs.append(log)
        self._persist(log)
        return log

    def get_all_logs(self) -> List[dict]:
        return [log.to_dict() for log in self._logs]

    def get_summary(self) -> dict:
        """Aggregate statistics across all logged requests."""
        if not self._logs:
            return {}

        total = len(self._logs)
        refused = sum(1 for l in self._logs if l.was_refused)
        latencies = [l.latency_ms for l in self._logs if not l.was_refused]
        avg_latency = sum(latencies) / len(latencies) if latencies else 0.0
        total_tokens = sum(
            l.estimated_prompt_tokens + l.estimated_response_tokens
            for l in self._logs
        )

        by_backend: dict = {}
        for log in self._logs:
            b = log.backend
            if b not in by_backend:
                by_backend[b] = {"count": 0, "latencies": [], "refused": 0}
            by_backend[b]["count"] += 1
            by_backend[b]["latencies"].append(log.latency_ms)
            if log.was_refused:
                by_backend[b]["refused"] += 1

        backend_summary = {}
        for b, data in by_backend.items():
            lats = data["latencies"]
            backend_summary[b] = {
                "requests": data["count"],
                "avg_latency_ms": round(sum(lats) / len(lats), 2) if lats else 0,
                "refusals": data["refused"],
            }

        return {
            "total_requests": total,
            "total_refusals": refused,
            "refusal_rate": round(refused / total, 3) if total else 0,
            "avg_latency_ms": round(avg_latency, 2),
            "total_estimated_tokens": total_tokens,
            "by_backend": backend_summary,
        }

    def new_session(self) -> None:
        """Start a new session (new session_id)."""
        self._session_id = str(uuid.uuid4())[:8]

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _ensure_log_dir(self) -> None:
        self.log_path.parent.mkdir(parents=True, exist_ok=True)

    def _load_existing_logs(self) -> None:
        """Load existing logs from disk on startup."""
        if not self.log_path.exists():
            return
        try:
            with open(self.log_path, "r") as f:
                raw = json.load(f)
            for entry in raw:
                try:
                    self._logs.append(RequestLog(**entry))
                except Exception:
                    pass  # Skip malformed entries
        except Exception as exc:
            logger.warning("Could not load existing logs: %s", exc)

    def _persist(self, log: RequestLog) -> None:
        """Append a single log entry to the JSON file."""
        try:
            existing: List[dict] = []
            if self.log_path.exists():
                with open(self.log_path, "r") as f:
                    existing = json.load(f)
            existing.append(log.to_dict())
            with open(self.log_path, "w") as f:
                json.dump(existing, f, indent=2)
        except Exception as exc:
            logger.error("Failed to persist log: %s", exc)


# ---------------------------------------------------------------------------
# Context manager for timing
# ---------------------------------------------------------------------------

class Timer:
    """Simple context manager for measuring elapsed time in milliseconds."""

    def __init__(self):
        self.elapsed_ms: float = 0.0
        self._start: float = 0.0

    def __enter__(self) -> "Timer":
        self._start = time.perf_counter()
        return self

    def __exit__(self, *_) -> None:
        self.elapsed_ms = (time.perf_counter() - self._start) * 1000