"""
FastAPI collector server for plyra-trace.

Routes:
  POST /v1/traces          — OTLP ingest (protobuf or JSON)
  GET  /api/v1/projects    — List projects
  GET  /api/v1/traces      — List traces with filters
  GET  /api/v1/traces/{id}/spans — Span tree for a trace
  GET  /api/v1/stats/llm   — LLM cost + token stats
  GET  /api/v1/stats/agents — Agent performance stats
  GET  /api/v1/stats/guards — Guard/policy audit events
  GET  /api/v1/stats/db    — Database metadata
  GET  /                   — UI (index.html)
"""

import asyncio
import json
import logging
import os
import sqlite3
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from datetime import UTC
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse

from plyra_trace.collector import storage
from plyra_trace.collector.models import (
    AgentsStatsResponse,
    DbMetaResponse,
    GuardsStatsResponse,
    LLMStatsResponse,
    ProjectsResponse,
    SpansResponse,
    TracesResponse,
)
from plyra_trace.collector.normalizer import SpanNormalizer

logger = logging.getLogger("plyra_trace.collector.server")

STATIC_DIR = Path(__file__).parent / "static"

# Global state (set on startup)
_db_conn: sqlite3.Connection | None = None
_db_path: str = "~/.plyra/traces.db"
_auth_enabled: bool = True

# SSE live-span subscribers (one asyncio.Queue per connected client)
_live_subscribers: list[asyncio.Queue] = []


async def broadcast_span(span: dict) -> None:
    """Push a normalised span dict to every connected SSE subscriber."""
    if not _live_subscribers:
        return
    data = json.dumps(span)
    for q in list(_live_subscribers):
        try:
            q.put_nowait(data)
        except asyncio.QueueFull:
            pass  # slow consumer — drop silently


def get_db() -> sqlite3.Connection:
    if _db_conn is None:
        raise RuntimeError("Database not initialized")
    return _db_conn


def _ns_to_iso(ns: int) -> str:
    """Convert nanoseconds-since-epoch to ISO 8601 string."""
    from datetime import datetime

    secs = ns / 1e9
    return datetime.fromtimestamp(secs, tz=UTC).isoformat()


def _parse_otlp_protobuf(body: bytes) -> list[dict[str, Any]]:
    """
    Parse an OTLP ExportTraceServiceRequest protobuf body.
    Returns a list of span dicts ready for storage.insert_span().

    Uses opentelemetry-proto if available, else falls back to JSON decode.
    """
    try:
        from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
            ExportTraceServiceRequest,
        )

        req = ExportTraceServiceRequest()
        req.ParseFromString(body)
        return _extract_spans_from_proto(req)
    except Exception as exc:
        logger.warning("Protobuf parse failed (%s), trying JSON fallback", exc)
        raise


def _extract_spans_from_proto(req: Any) -> list[dict[str, Any]]:
    """Extract span dicts from a parsed protobuf ExportTraceServiceRequest."""
    spans = []
    for resource_span in req.resource_spans:
        # Extract service.name from resource attributes
        project_id = "unknown"
        for attr in resource_span.resource.attributes:
            if attr.key == "service.name":
                project_id = attr.value.string_value or "unknown"

        for scope_span in resource_span.scope_spans:
            for span in scope_span.spans:
                span_id = span.span_id.hex()
                trace_id = span.trace_id.hex()
                parent_id = span.parent_span_id.hex() if span.parent_span_id else None
                if parent_id and all(c == "0" for c in parent_id):
                    parent_id = None

                attrs: dict[str, Any] = {}
                for attr in span.attributes:
                    val = attr.value
                    if val.HasField("string_value"):
                        attrs[attr.key] = val.string_value
                    elif val.HasField("int_value"):
                        attrs[attr.key] = val.int_value
                    elif val.HasField("double_value"):
                        attrs[attr.key] = val.double_value
                    elif val.HasField("bool_value"):
                        attrs[attr.key] = val.bool_value

                events = []
                for evt in span.events:
                    evt_attrs: dict[str, Any] = {}
                    for attr in evt.attributes:
                        val = attr.value
                        if val.HasField("string_value"):
                            evt_attrs[attr.key] = val.string_value
                        elif val.HasField("int_value"):
                            evt_attrs[attr.key] = val.int_value
                    events.append(
                        {
                            "name": evt.name,
                            "timestamp": _ns_to_iso(evt.time_unix_nano),
                            "attributes": evt_attrs,
                        }
                    )

                kind_map = {
                    0: "CHAIN",
                    1: "CHAIN",
                    2: "CHAIN",
                    3: "LLM",
                    4: "TOOL",
                    5: "CHAIN",
                }
                plyra_kind = attrs.get("plyra.span.kind") or attrs.get("openinference.span.kind")
                if not plyra_kind:
                    plyra_kind = kind_map.get(span.kind, "CHAIN")

                started_at = _ns_to_iso(span.start_time_unix_nano)
                ended_at = _ns_to_iso(span.end_time_unix_nano)
                duration_ms = max(0, int((span.end_time_unix_nano - span.start_time_unix_nano) / 1e6))

                status = "ok"
                if span.status.code == 2:  # STATUS_CODE_ERROR
                    status = "error"

                spans.append(
                    {
                        "span_id": span_id,
                        "trace_id": trace_id,
                        "parent_id": parent_id,
                        "project_id": project_id,
                        "name": span.name,
                        "kind": plyra_kind,
                        "started_at": started_at,
                        "ended_at": ended_at,
                        "duration_ms": duration_ms,
                        "attributes": attrs,
                        "events": events,
                        "status": status,
                    }
                )
    return spans


def _parse_otlp_json(body: bytes) -> list[dict[str, Any]]:
    """
    Parse OTLP JSON (application/json content-type).
    Handles both OpenTelemetry JSON format and plyra-trace simplified format.
    """
    data = json.loads(body)
    spans = []

    # Standard OTLP JSON format
    resource_spans = data.get("resourceSpans") or data.get("resource_spans", [])
    for rs in resource_spans:
        project_id = "unknown"
        resource = rs.get("resource", {})
        for attr in resource.get("attributes", []):
            if attr.get("key") == "service.name":
                project_id = (
                    attr.get("value", {}).get("stringValue") or attr.get("value", {}).get("string_value") or "unknown"
                )

        scope_spans_list = rs.get("scopeSpans") or rs.get("scope_spans", [])
        for ss in scope_spans_list:
            for span in ss.get("spans", []):
                attrs: dict[str, Any] = {}
                for attr in span.get("attributes", []):
                    val = attr.get("value", {})
                    if "stringValue" in val:
                        attrs[attr["key"]] = val["stringValue"]
                    elif "string_value" in val:
                        attrs[attr["key"]] = val["string_value"]
                    elif "intValue" in val:
                        attrs[attr["key"]] = int(val["intValue"])
                    elif "int_value" in val:
                        attrs[attr["key"]] = int(val["int_value"])
                    elif "doubleValue" in val:
                        attrs[attr["key"]] = float(val["doubleValue"])
                    elif "double_value" in val:
                        attrs[attr["key"]] = float(val["double_value"])
                    elif "boolValue" in val:
                        attrs[attr["key"]] = bool(val["boolValue"])
                    elif "bool_value" in val:
                        attrs[attr["key"]] = bool(val["bool_value"])

                span_id = span.get("spanId") or span.get("span_id", "")
                trace_id = span.get("traceId") or span.get("trace_id", "")
                parent_id = span.get("parentSpanId") or span.get("parent_span_id") or None

                plyra_kind = attrs.get("plyra.span.kind") or attrs.get("openinference.span.kind") or "CHAIN"

                start_ns = int(span.get("startTimeUnixNano") or span.get("start_time_unix_nano") or 0)
                end_ns = int(span.get("endTimeUnixNano") or span.get("end_time_unix_nano") or 0)
                started_at = _ns_to_iso(start_ns) if start_ns else ""
                ended_at = _ns_to_iso(end_ns) if end_ns else ""
                duration_ms = max(0, int((end_ns - start_ns) / 1e6)) if start_ns and end_ns else 0

                # Fallback: accept simplified plyra-trace dicts directly
                if not started_at and "started_at" in span:
                    started_at = span["started_at"]
                    ended_at = span.get("ended_at", started_at)
                    duration_ms = span.get("duration_ms", 0)

                status_data = span.get("status", {})
                status = "ok"
                if isinstance(status_data, dict):
                    code = status_data.get("code") or status_data.get("code", 0)
                    if code == 2 or str(code).upper() == "STATUS_CODE_ERROR":
                        status = "error"

                events = []
                for evt in span.get("events", []):
                    evt_attrs: dict[str, Any] = {}
                    for attr in evt.get("attributes", []):
                        val = attr.get("value", {})
                        if "stringValue" in val:
                            evt_attrs[attr["key"]] = val["stringValue"]
                    events.append(
                        {
                            "name": evt.get("name", ""),
                            "timestamp": _ns_to_iso(int(evt.get("timeUnixNano") or evt.get("time_unix_nano") or 0)),
                            "attributes": evt_attrs,
                        }
                    )

                spans.append(
                    {
                        "span_id": span_id,
                        "trace_id": trace_id,
                        "parent_id": parent_id or None,
                        "project_id": project_id,
                        "name": span.get("name", ""),
                        "kind": plyra_kind,
                        "started_at": started_at,
                        "ended_at": ended_at,
                        "duration_ms": duration_ms,
                        "attributes": attrs,
                        "events": events,
                        "status": status,
                    }
                )

    # Simplified format: {"spans": [...]} where each span is a plyra-style dict
    if not spans and "spans" in data:
        for span in data["spans"]:
            spans.append(
                {
                    "span_id": span.get("span_id", ""),
                    "trace_id": span.get("trace_id", ""),
                    "parent_id": span.get("parent_id"),
                    "project_id": span.get("project_id", data.get("project", "unknown")),
                    "name": span.get("name", ""),
                    "kind": span.get("kind", "CHAIN"),
                    "started_at": span.get("started_at", ""),
                    "ended_at": span.get("ended_at", ""),
                    "duration_ms": span.get("duration_ms", 0),
                    "attributes": span.get("attributes", {}),
                    "events": span.get("events", []),
                    "status": span.get("status", "ok"),
                }
            )

    return spans


# ── FastAPI app ────────────────────────────────────────────────────────────────


def create_app(auth_enabled: bool = True) -> FastAPI:
    """Create the FastAPI application."""

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
        yield

    app = FastAPI(
        title="plyra-trace collector",
        version="2.0.0",
        lifespan=lifespan,
    )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # ── OTLP ingest ────────────────────────────────────────────────────────────

    @app.post("/v1/traces")
    async def ingest_traces(request: Request) -> JSONResponse:
        """
        Accept OTLP spans (protobuf or JSON).
        Parses, stores each span, and updates trace summaries.
        """
        db = get_db()
        content_type = request.headers.get("content-type", "")
        body = await request.body()

        # Auth check
        if auth_enabled:
            auth_header = request.headers.get("authorization", "")
            if auth_header.startswith("Bearer "):
                key = auth_header[7:]
                from plyra_trace.collector.auth import validate_api_key

                if not validate_api_key(key):
                    raise HTTPException(status_code=401, detail="Invalid or revoked API key")
            else:
                raise HTTPException(status_code=401, detail="Authorization header required")

        # Parse spans
        spans: list[dict[str, Any]] = []
        if "protobuf" in content_type or "application/x-protobuf" in content_type:
            try:
                spans = _parse_otlp_protobuf(body)
            except Exception as exc:
                logger.error("Failed to parse protobuf: %s", exc)
                raise HTTPException(status_code=400, detail=f"Protobuf parse error: {exc}")
        else:
            try:
                spans = _parse_otlp_json(body)
            except Exception as exc:
                logger.error("Failed to parse JSON OTLP: %s", exc)
                raise HTTPException(status_code=400, detail=f"JSON parse error: {exc}")

        if not spans:
            return JSONResponse(content={})

        # Normalize spans from any source schema
        normalizer = SpanNormalizer()
        spans = [normalizer.normalize(s) for s in spans]

        # Upsert projects and spans
        seen_projects: set[str] = set()
        seen_traces: set[str] = set()

        for span in spans:
            pid = span.get("project_id", "unknown")
            if pid not in seen_projects:
                storage.upsert_project(db, pid)
                seen_projects.add(pid)
            storage.insert_span(db, span)
            await broadcast_span(span)
            seen_traces.add(span["trace_id"])

        # Update trace summaries
        for trace_id in seen_traces:
            storage.update_trace_summary(db, trace_id)

        logger.debug("Ingested %d spans across %d traces", len(spans), len(seen_traces))
        return JSONResponse(content={})

    # ── REST API ───────────────────────────────────────────────────────────────

    @app.get("/api/v1/projects", response_model=ProjectsResponse)
    async def get_projects_endpoint() -> dict[str, Any]:
        """List all projects with trace statistics."""
        db = get_db()
        projects = storage.get_projects_with_stats(db)
        return {"projects": projects}

    @app.get("/api/v1/traces", response_model=TracesResponse)
    async def get_traces_endpoint(
        project: str | None = Query(None),
        limit: int = Query(50, ge=1, le=500),
        offset: int = Query(0, ge=0),
        since: str | None = Query(None, description="e.g. '24h', '7d'"),
    ) -> dict[str, Any]:
        """List trace summaries with optional filters."""
        db = get_db()
        since_hours = _parse_since(since)
        traces, total = storage.get_traces(
            db,
            project_id=project,
            limit=limit,
            offset=offset,
            since_hours=since_hours,
        )
        # Coerce booleans for Pydantic
        for t in traces:
            t["has_guard_block"] = bool(t.get("has_guard_block", 0))
            t["has_error"] = bool(t.get("has_error", 0))
        return {"traces": traces, "total": total}

    @app.get("/api/v1/traces/{trace_id}/spans", response_model=SpansResponse)
    async def get_spans_endpoint(trace_id: str) -> dict[str, Any]:
        """Retrieve all spans for a specific trace."""
        db = get_db()
        spans = storage.get_spans(db, trace_id)
        if not spans:
            raise HTTPException(status_code=404, detail="Trace not found")
        return {"spans": spans}

    @app.get("/api/v1/stats/llm", response_model=LLMStatsResponse)
    async def get_stats_llm_endpoint(
        project: str | None = Query(None),
        since: str | None = Query("24h"),
    ) -> dict[str, Any]:
        """LLM cost, token, and cache statistics."""
        db = get_db()
        since_hours = _parse_since(since) or 24
        return storage.get_stats_llm(db, project_id=project, since_hours=since_hours)

    @app.get("/api/v1/stats/agents", response_model=AgentsStatsResponse)
    async def get_stats_agents_endpoint(
        project: str | None = Query(None),
        since: str | None = Query("24h"),
    ) -> dict[str, Any]:
        """Agent performance statistics."""
        db = get_db()
        since_hours = _parse_since(since) or 24
        agents = storage.get_stats_agents(db, project_id=project, since_hours=since_hours)
        return {"agents": agents}

    @app.get("/api/v1/stats/guards", response_model=GuardsStatsResponse)
    async def get_stats_guards_endpoint(
        project: str | None = Query(None),
        since: str | None = Query("24h"),
    ) -> dict[str, Any]:
        """Guard/policy audit events."""
        db = get_db()
        since_hours = _parse_since(since) or 24
        return storage.get_stats_guards(db, project_id=project, since_hours=since_hours)

    @app.get("/api/v1/stats/db", response_model=DbMetaResponse)
    async def get_stats_db_endpoint() -> dict[str, Any]:
        """Database metadata."""
        db = get_db()
        return storage.get_db_meta(db, _db_path)

    @app.get("/api/v1/live/spans")
    async def live_spans_sse(request: Request, project: str | None = Query(None)) -> Any:
        """SSE stream: emits a JSON span object for every span ingested after connection."""
        from sse_starlette.sse import EventSourceResponse

        queue: asyncio.Queue = asyncio.Queue(maxsize=500)
        _live_subscribers.append(queue)

        async def event_generator():
            try:
                while True:
                    if await request.is_disconnected():
                        break
                    try:
                        data = await asyncio.wait_for(queue.get(), timeout=15.0)
                        span = json.loads(data)
                        if project is None or span.get("project_id") == project:
                            yield {"data": data}
                    except TimeoutError:
                        yield {"data": '{"type":"keepalive"}'}
            finally:
                if queue in _live_subscribers:
                    _live_subscribers.remove(queue)

        return EventSourceResponse(event_generator())

    # ── Static UI ──────────────────────────────────────────────────────────────
    # Serve index.html for GET / (must be mounted AFTER API routes)
    if STATIC_DIR.exists():

        @app.get("/")
        async def serve_index() -> FileResponse:
            """Serve the plyra-trace UI."""
            return FileResponse(STATIC_DIR / "index.html")

    return app


def _parse_since(since: str | None) -> int | None:
    """Parse a since string like '24h' or '7d' into hours."""
    if not since:
        return None
    since = since.strip().lower()
    if since.endswith("h"):
        try:
            return int(since[:-1])
        except ValueError:
            return None
    if since.endswith("d"):
        try:
            return int(since[:-1]) * 24
        except ValueError:
            return None
    return None


def run_server(
    host: str = "127.0.0.1",
    port: int = 7700,
    otlp_port: int | None = None,
    db_path: str = "~/.plyra/traces.db",
    auth_enabled: bool = True,
) -> None:
    """
    Start the plyra-trace collector server.

    Runs a single uvicorn server:
    - UI + REST API: http://{host}:{port}
    - OTLP ingest:   http://{host}:{port}/v1/traces

    Args:
        host: Bind address (default: 127.0.0.1)
        port: Port for UI + API + OTLP (default: 7700)
        otlp_port: Ignored (reserved for future dual-port mode)
        db_path: SQLite database path (default: ~/.plyra/traces.db)
        auth_enabled: Require API key authentication

    Example:
        >>> from plyra_trace.collector.server import run_server
        >>> run_server(host="0.0.0.0", port=7700, auth_enabled=False)
    """
    global _db_conn, _db_path, _auth_enabled

    import uvicorn

    from plyra_trace import __version__

    _db_path = db_path
    _auth_enabled = auth_enabled
    expanded_db = os.path.expanduser(db_path)
    _db_conn = storage.open_db(db_path)

    db_meta = storage.get_db_meta(_db_conn, db_path)

    trace_note = f"({db_meta['trace_count']} traces)"
    auth_note = (
        "disabled  ·  no key required"
        if not auth_enabled
        else ("enabled  ·  create key: plyra-trace keys create --project NAME")
    )

    _banner = f"""
  ◈ plyra-trace v{__version__}
  {"─" * 46}
  UI:         http://{host}:{port}
  OTLP:       http://{host}:{port}/v1/traces
  Database:   {expanded_db}  {trace_note}
  Auth:       {auth_note}
  {"─" * 46}

  In your agent:
      import plyra_trace
      pt = plyra_trace.init(
          project="my-agent",
          endpoint="http://{host}:{port}",
      )
  {"─" * 46}
"""
    print(_banner)

    app = create_app(auth_enabled=auth_enabled)

    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="warning",
        access_log=False,
    )
