"""Trace ingestion, querying, and monitoring routes."""

from __future__ import annotations

import json
import logging
import math
import os
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel as PydanticBaseModel
from pydantic import Field
from sqlalchemy import func
from sqlalchemy.orm import Session

from evalkit.models import Suite, TestCase
from evalkit.server.db.models import ProjectModel, TraceModel, TraceScoreModel
from evalkit.server.db.session import SessionLocal, get_db
from evalkit.suites import BUILTIN_SUITES

logger = logging.getLogger("evalkit.server")
templates = Jinja2Templates(directory=str(Path(__file__).parent.parent / "templates"))

# -- Request / Response models -------------------------------------------------


class TraceInput(PydanticBaseModel):
    test_case: TestCase
    timestamp: str | None = None
    metadata: dict[str, str] = {}


class IngestRequest(PydanticBaseModel):
    project_id: int
    suite: Suite | None = None
    judge_provider: str | None = None
    judge_model: str | None = None
    api_key: str | None = None
    traces: list[TraceInput] = Field(max_length=10000)


class IngestResponse(PydanticBaseModel):
    count: int
    trace_ids: list[int]


# -- Helpers -------------------------------------------------------------------


def _trace_to_dict(trace: TraceModel) -> dict[str, Any]:
    scores = [
        {
            "metric_name": s.metric_name,
            "score": s.score,
            "passed": bool(s.passed),
            "reason": s.reason,
        }
        for s in trace.scores
    ]
    avg = sum(s.score for s in trace.scores) / len(trace.scores) if trace.scores else None
    tc = json.loads(trace.test_case_json)
    return {
        "id": trace.id,
        "project_id": trace.project_id,
        "timestamp": trace.timestamp.isoformat(),
        "test_case": tc,
        "question": tc.get("question", ""),
        "answer": tc.get("answer", ""),
        "context": tc.get("context", []),
        "expected_answer": tc.get("expected_answer"),
        "scores": scores,
        "metadata": trace.get_metadata(),
        "avg_score": round(avg, 4) if avg is not None else None,
    }


def _parse_iso(value: str) -> datetime:
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


# -- Background evaluation ----------------------------------------------------


async def _evaluate_traces(trace_ids: list[int], request: IngestRequest):
    from evalkit.engine import EvaluationEngine

    suite_key = request.suite or Suite(_get_default("EVALKIT_SUITE", Suite.RAG_QA.value))
    suite_factory = BUILTIN_SUITES.get(suite_key)
    if suite_factory is None:
        logger.error("Unknown suite for trace evaluation: %s", suite_key)
        return
    suite_instance = suite_factory()

    judge = None
    if suite_instance.requires_judge:
        from evalkit.judges.anthropic import AnthropicJudge
        from evalkit.judges.ollama import OllamaJudge
        from evalkit.judges.openai import AzureOpenAIJudge, OpenAIJudge

        provider = request.judge_provider or _get_default("EVALKIT_JUDGE_PROVIDER", "openai")
        model = request.judge_model or os.environ.get("EVALKIT_JUDGE_MODEL")
        if provider == "openai":
            judge = OpenAIJudge(model=model or "gpt-4o-mini", api_key=request.api_key)
        elif provider == "azure":
            azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
            judge = AzureOpenAIJudge(
                model=model or "gpt-4o-mini",
                api_key=request.api_key,
                azure_endpoint=azure_endpoint,
            )
        elif provider == "anthropic":
            judge = AnthropicJudge(model=model or "claude-haiku-4-5-20251001", api_key=request.api_key)
        else:
            ollama_url = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
            judge = OllamaJudge(model=model or "gpt-oss", base_url=ollama_url)

    engine = EvaluationEngine(judge=judge)

    db = SessionLocal()
    try:
        traces = [db.get(TraceModel, tid) for tid in trace_ids]
        test_cases = [TestCase.model_validate_json(t.test_case_json) for t in traces]

        report = await engine.run_async(test_cases=test_cases, suite=suite_instance)

        for trace, result in zip(traces, report.results):
            for ms in result.scores:
                db.add(
                    TraceScoreModel(
                        trace_id=trace.id,
                        metric_name=ms.metric_name,
                        score=ms.score,
                        passed=int(ms.passed),
                        reason=ms.reason,
                    )
                )
        db.commit()
        logger.info("Trace evaluation complete: %d traces scored with suite=%s", len(trace_ids), suite_key)
    except Exception:
        logger.exception("Trace evaluation failed for trace_ids=%s", trace_ids)
    finally:
        db.close()


def _get_default(key: str, fallback: str) -> str:
    return os.environ.get(key, fallback)


# -- API routes (mounted under /api/v1) ----------------------------------------

api_router = APIRouter()


@api_router.post("/traces", response_model=IngestResponse, status_code=201)
async def ingest_traces(
    request: IngestRequest,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
):
    project = db.get(ProjectModel, request.project_id)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")

    trace_ids: list[int] = []
    for t in request.traces:
        ts = _parse_iso(t.timestamp) if t.timestamp else datetime.now(UTC)
        trace = TraceModel(
            project_id=request.project_id,
            timestamp=ts,
            test_case_json=t.test_case.model_dump_json(),
            metadata_json=json.dumps(t.metadata),
        )
        db.add(trace)
        db.flush()
        trace_ids.append(trace.id)

    db.commit()

    # Evaluate in background if suite specified (or server default exists)
    background_tasks.add_task(_evaluate_traces, trace_ids, request)

    return IngestResponse(count=len(trace_ids), trace_ids=trace_ids)


@api_router.get("/traces")
def list_traces(
    project_id: int,
    start: str | None = None,
    end: str | None = None,
    limit: int = Query(50, ge=1, le=1000),
    offset: int = Query(0, ge=0),
    db: Session = Depends(get_db),
):
    project = db.get(ProjectModel, project_id)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")

    query = db.query(TraceModel).filter(TraceModel.project_id == project_id)
    count_query = db.query(func.count(TraceModel.id)).filter(TraceModel.project_id == project_id)

    if start:
        ts = _parse_iso(start)
        query = query.filter(TraceModel.timestamp >= ts)
        count_query = count_query.filter(TraceModel.timestamp >= ts)
    if end:
        ts = _parse_iso(end)
        query = query.filter(TraceModel.timestamp <= ts)
        count_query = count_query.filter(TraceModel.timestamp <= ts)

    total = count_query.scalar()
    traces = query.order_by(TraceModel.timestamp.desc()).offset(offset).limit(limit).all()
    return {"traces": [_trace_to_dict(t) for t in traces], "total": total}


@api_router.get("/traces/{trace_id}")
def get_trace(trace_id: int, db: Session = Depends(get_db)):
    trace = db.get(TraceModel, trace_id)
    if trace is None:
        raise HTTPException(status_code=404, detail="Trace not found")
    return _trace_to_dict(trace)


@api_router.get("/projects/{project_id}/monitor")
def get_monitor_data(
    project_id: int,
    start: str | None = None,
    end: str | None = None,
    interval: str = "1h",
    metric: str | None = None,
    db: Session = Depends(get_db),
):
    project = db.get(ProjectModel, project_id)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")

    valid_intervals = {"5m": 300, "15m": 900, "1h": 3600, "6h": 21600, "1d": 86400}
    interval_seconds = valid_intervals.get(interval)
    if interval_seconds is None:
        raise HTTPException(status_code=400, detail=f"interval must be one of {list(valid_intervals.keys())}")

    query = db.query(TraceModel).filter(TraceModel.project_id == project_id)
    if start:
        query = query.filter(TraceModel.timestamp >= _parse_iso(start))
    if end:
        query = query.filter(TraceModel.timestamp <= _parse_iso(end))

    traces = query.order_by(TraceModel.timestamp.asc()).all()

    buckets_map: dict[str, list[float]] = {}
    for trace in traces:
        scores = trace.scores
        if metric:
            scores = [s for s in scores if s.metric_name == metric]
        if not scores:
            continue
        avg = sum(s.score for s in scores) / len(scores)

        epoch = trace.timestamp.timestamp()
        bucket_epoch = (int(epoch) // interval_seconds) * interval_seconds
        key = datetime.fromtimestamp(bucket_epoch, tz=UTC).isoformat()
        buckets_map.setdefault(key, []).append(avg)

    buckets = []
    for key in sorted(buckets_map.keys()):
        values = buckets_map[key]
        buckets.append(
            {
                "timestamp": key,
                "count": len(values),
                "avg_score": round(sum(values) / len(values), 4),
                "min_score": round(min(values), 4),
                "max_score": round(max(values), 4),
                "anomaly": False,
            }
        )

    if len(buckets) >= 3:
        all_avgs = [b["avg_score"] for b in buckets]
        overall_mean = sum(all_avgs) / len(all_avgs)
        variance = sum((x - overall_mean) ** 2 for x in all_avgs) / len(all_avgs)
        stddev = math.sqrt(variance)
        threshold = overall_mean - 2 * stddev

        for i, bucket in enumerate(buckets):
            if bucket["avg_score"] < threshold:
                bucket["anomaly"] = True
            if i > 0 and buckets[i - 1]["avg_score"] > 0:
                drop = (buckets[i - 1]["avg_score"] - bucket["avg_score"]) / buckets[i - 1]["avg_score"]
                if drop > 0.15:
                    bucket["anomaly"] = True

    total_traces = sum(b["count"] for b in buckets)
    all_scores = [b["avg_score"] for b in buckets]
    anomaly_count = sum(1 for b in buckets if b["anomaly"])

    return {
        "buckets": buckets,
        "summary": {
            "total_traces": total_traces,
            "avg_score": round(sum(all_scores) / len(all_scores), 4) if all_scores else 0,
            "anomaly_count": anomaly_count,
        },
    }


# -- HTML routes (mounted at root) ---------------------------------------------

html_router = APIRouter()

TRACES_PER_PAGE_OPTIONS = [5, 10, 20, 50]
DEFAULT_PER_PAGE = 50

INTERVAL_DELTAS = {
    "5m": timedelta(minutes=5),
    "15m": timedelta(minutes=15),
    "1h": timedelta(hours=1),
    "6h": timedelta(hours=6),
    "1d": timedelta(days=1),
}

RANGE_DELTAS = {
    "1h": timedelta(hours=1),
    "6h": timedelta(hours=6),
    "24h": timedelta(hours=24),
    "7d": timedelta(days=7),
    "30d": timedelta(days=30),
}


def _bucket_key(ts: datetime, interval_seconds: int) -> str:
    epoch = ts.timestamp()
    bucket_epoch = (int(epoch) // interval_seconds) * interval_seconds

    dt = datetime.fromtimestamp(bucket_epoch, tz=UTC)
    if interval_seconds < 3600:
        return dt.strftime("%H:%M")
    if interval_seconds < 86400:
        return dt.strftime("%m/%d %H:%M")
    return dt.strftime("%m/%d")


@html_router.get("/projects/{project_id}/monitor", response_class=HTMLResponse)
def monitor_page(
    project_id: int,
    request: Request,
    page: int = 1,
    per_page: int = DEFAULT_PER_PAGE,
    range: str = "24h",
    interval: str = "1h",
    agg: str = "avg",
    metric: str | None = None,
    db: Session = Depends(get_db),
):
    project = db.get(ProjectModel, project_id)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")

    if per_page not in TRACES_PER_PAGE_OPTIONS:
        per_page = DEFAULT_PER_PAGE

    now = datetime.now(UTC)
    delta = RANGE_DELTAS.get(range, timedelta(hours=24))
    start = now - delta
    interval_delta = INTERVAL_DELTAS.get(interval, timedelta(hours=1))
    interval_seconds = int(interval_delta.total_seconds())

    if page < 1:
        page = 1
    trace_query = (
        db.query(TraceModel)
        .filter(TraceModel.project_id == project_id, TraceModel.timestamp >= start)
        .order_by(TraceModel.timestamp.desc())
    )
    total = trace_query.count()
    total_pages = max(1, (total + per_page - 1) // per_page)
    if page > total_pages:
        page = total_pages
    traces = trace_query.offset((page - 1) * per_page).limit(per_page).all()

    # Load all traces in range for chart
    all_traces = (
        db.query(TraceModel)
        .filter(TraceModel.project_id == project_id, TraceModel.timestamp >= start)
        .order_by(TraceModel.timestamp.asc())
        .all()
    )

    # Build buckets with per-trace avg scores
    buckets_map: dict[str, list[float]] = {}
    for trace in all_traces:
        trace_scores = trace.scores
        if metric:
            trace_scores = [s for s in trace_scores if s.metric_name == metric]
        if not trace_scores:
            continue
        trace_avg = sum(s.score for s in trace_scores) / len(trace_scores)
        key = _bucket_key(trace.timestamp, interval_seconds)
        buckets_map.setdefault(key, []).append(trace_avg)

    chart_labels = list(sorted(buckets_map.keys()))
    chart_counts = [len(buckets_map[k]) for k in chart_labels]

    # Compute score per bucket based on selected aggregation
    agg_fn = {"avg": lambda v: sum(v) / len(v), "min": min, "max": max}.get(agg, lambda v: sum(v) / len(v))
    chart_scores = [round(agg_fn(buckets_map[k]), 4) for k in chart_labels]

    # Bucket-level anomaly detection (when multiple buckets exist)
    chart_anomalies = [False] * len(chart_scores)
    if len(chart_scores) >= 3:
        mean = sum(chart_scores) / len(chart_scores)
        variance = sum((x - mean) ** 2 for x in chart_scores) / len(chart_scores)
        stddev = math.sqrt(variance)
        threshold = mean - 2 * stddev
        for i, score in enumerate(chart_scores):
            if score < threshold:
                chart_anomalies[i] = True
            if i > 0 and chart_scores[i - 1] > 0:
                drop = (chart_scores[i - 1] - score) / chart_scores[i - 1]
                if drop > 0.15:
                    chart_anomalies[i] = True

    # Per-trace anomaly detection: flag individual traces whose avg score
    # is more than 2 stddev below the overall mean of all trace scores.
    all_trace_avgs = []
    for vals in buckets_map.values():
        all_trace_avgs.extend(vals)

    trace_anomaly_ids: set[int] = set()
    if len(all_trace_avgs) >= 2:
        overall_mean = sum(all_trace_avgs) / len(all_trace_avgs)
        overall_var = sum((x - overall_mean) ** 2 for x in all_trace_avgs) / len(all_trace_avgs)
        overall_stddev = math.sqrt(overall_var)
        trace_threshold = overall_mean - 2 * overall_stddev
        for trace in all_traces:
            trace_scores_list = trace.scores
            if metric:
                trace_scores_list = [s for s in trace_scores_list if s.metric_name == metric]
            if not trace_scores_list:
                continue
            tavg = sum(s.score for s in trace_scores_list) / len(trace_scores_list)
            if tavg < trace_threshold:
                trace_anomaly_ids.add(trace.id)

    anomaly_count = sum(chart_anomalies) + len(trace_anomaly_ids)

    # Unique metric names for filter dropdown
    metric_names = (
        db.query(TraceScoreModel.metric_name)
        .join(TraceModel)
        .filter(TraceModel.project_id == project_id)
        .distinct()
        .all()
    )
    metric_names = sorted({m[0] for m in metric_names})

    return templates.TemplateResponse(
        request,
        "monitor.html",
        {
            "project": project,
            "traces": traces,
            "total": total,
            "page": page,
            "total_pages": total_pages,
            "time_range": range,
            "interval": interval,
            "agg": agg,
            "metric": metric,
            "metric_names": metric_names,
            "chart_labels": chart_labels,
            "chart_scores": chart_scores,
            "chart_counts": chart_counts,
            "chart_anomalies": chart_anomalies,
            "anomaly_count": anomaly_count,
            "trace_anomaly_ids": trace_anomaly_ids,
            "per_page": per_page,
            "per_page_options": TRACES_PER_PAGE_OPTIONS,
        },
    )


@html_router.get("/traces/{trace_id}", response_class=HTMLResponse)
def trace_detail_page(trace_id: int, request: Request, db: Session = Depends(get_db)):
    trace = db.get(TraceModel, trace_id)
    if trace is None:
        raise HTTPException(status_code=404, detail="Trace not found")
    tc = json.loads(trace.test_case_json)
    return templates.TemplateResponse(
        request,
        "trace.html",
        {
            "trace": trace,
            "project": trace.project,
            "test_case": tc,
            "context": tc.get("context", []),
            "metadata": trace.get_metadata(),
        },
    )
