"""Experiment management API endpoints."""
from __future__ import annotations

import csv
import io
import json

from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from api.deps import (
    _effective_identity,
    _identity_user_id,
    _require_scope,
    _record_audit,
)
from api.models import ExperimentCreateRequest
from api.state import ROLE_LEVEL

from src.user_store import (
    create_experiment,
    delete_experiment,
    get_experiment,
    list_experiment_runs,
    list_experiments,
    update_experiment,
)

router = APIRouter()


class ExperimentUpdateRequest(BaseModel):
    name: str | None = Field(default=None, min_length=1, max_length=200)
    description: str | None = Field(default=None, max_length=1000)
    config: dict | None = None


def _is_admin(request: Request) -> bool:
    identity = _effective_identity(request)
    role = str(identity.get("role") or "viewer")
    return ROLE_LEVEL.get(role, 0) >= ROLE_LEVEL.get("owner", 4)


def _check_experiment_access(request: Request, exp: dict) -> None:
    """Raise 403 if non-admin user doesn't own the experiment."""
    if not _is_admin(request):
        caller = _identity_user_id(request) or ""
        if exp.get("user_id") != caller:
            raise HTTPException(status_code=403, detail="Access denied")


# ---------------------------------------------------------------------------
# GET /api/experiments
# ---------------------------------------------------------------------------

@router.get("/api/experiments")
def api_list_experiments(
    request: Request,
    experiment_type: str | None = Query(default=None),
    status: str | None = Query(default=None),
    limit: int = Query(default=50, ge=1, le=200),
    offset: int = Query(default=0, ge=0),
):
    _require_scope(request, "experiments:read")
    caller_user_id = _identity_user_id(request) or ""
    is_admin = _is_admin(request)

    items, total = list_experiments(
        user_id=caller_user_id,
        limit=limit,
        offset=offset,
        experiment_type=experiment_type,
        status=status,
        include_all_users=is_admin,
    )

    return {
        "ok": True,
        "data": {"items": items, "total": total, "limit": limit, "offset": offset},
    }


# ---------------------------------------------------------------------------
# POST /api/experiments
# ---------------------------------------------------------------------------

@router.post("/api/experiments")
def api_create_experiment(payload: ExperimentCreateRequest, request: Request):
    _require_scope(request, "experiments:write")
    caller_user_id = _identity_user_id(request) or ""

    exp_id = create_experiment(
        user_id=caller_user_id,
        name=payload.name,
        experiment_type=payload.experiment_type,
        config=payload.config,
        description=payload.description,
        project_id=payload.project_id,
    )

    _record_audit(
        request,
        action="experiment.create",
        resource_type="experiment",
        resource_id=str(exp_id),
        detail={"name": payload.name, "type": payload.experiment_type},
    )

    exp = get_experiment(exp_id)
    return {"ok": True, "data": exp}


# ---------------------------------------------------------------------------
# GET /api/experiments/{exp_id}
# ---------------------------------------------------------------------------

@router.get("/api/experiments/{exp_id}")
def api_get_experiment(exp_id: int, request: Request):
    _require_scope(request, "experiments:read")

    exp = get_experiment(exp_id)
    if not exp:
        raise HTTPException(status_code=404, detail="Experiment not found")

    _check_experiment_access(request, exp)
    return {"ok": True, "data": exp}


# ---------------------------------------------------------------------------
# PATCH /api/experiments/{exp_id}
# ---------------------------------------------------------------------------

@router.patch("/api/experiments/{exp_id}")
def api_update_experiment(exp_id: int, payload: ExperimentUpdateRequest, request: Request):
    _require_scope(request, "experiments:write")
    caller_user_id = _identity_user_id(request) or ""

    exp = get_experiment(exp_id)
    if not exp:
        raise HTTPException(status_code=404, detail="Experiment not found")

    _check_experiment_access(request, exp)

    owner_id = caller_user_id if not _is_admin(request) else None
    updated = update_experiment(
        exp_id,
        user_id=owner_id,
        name=payload.name,
        description=payload.description,
        config=payload.config,
    )
    if not updated:
        raise HTTPException(status_code=400, detail="No changes applied")

    _record_audit(
        request,
        action="experiment.update",
        resource_type="experiment",
        resource_id=str(exp_id),
        detail={"name": payload.name},
    )

    exp = get_experiment(exp_id)
    return {"ok": True, "data": exp}


# ---------------------------------------------------------------------------
# DELETE /api/experiments/{exp_id}
# ---------------------------------------------------------------------------

@router.delete("/api/experiments/{exp_id}")
def api_delete_experiment(exp_id: int, request: Request):
    _require_scope(request, "experiments:write")
    caller_user_id = _identity_user_id(request) or ""

    exp = get_experiment(exp_id)
    if not exp:
        raise HTTPException(status_code=404, detail="Experiment not found")

    is_admin = _is_admin(request)
    _check_experiment_access(request, exp)

    deleted = delete_experiment(exp_id, caller_user_id if not is_admin else exp.get("user_id", ""))
    if not deleted:
        raise HTTPException(status_code=500, detail="Failed to delete experiment")

    _record_audit(
        request,
        action="experiment.delete",
        resource_type="experiment",
        resource_id=str(exp_id),
        detail={"name": exp.get("name"), "type": exp.get("experiment_type")},
    )

    return {"ok": True, "data": {"deleted": True, "experiment_id": exp_id}}


# ---------------------------------------------------------------------------
# GET /api/experiments/{exp_id}/runs
# ---------------------------------------------------------------------------

@router.get("/api/experiments/{exp_id}/runs")
def api_list_experiment_runs(
    exp_id: int,
    request: Request,
    limit: int = Query(default=100, ge=1, le=500),
    offset: int = Query(default=0, ge=0),
):
    _require_scope(request, "experiments:read")

    exp = get_experiment(exp_id)
    if not exp:
        raise HTTPException(status_code=404, detail="Experiment not found")

    _check_experiment_access(request, exp)

    runs, total = list_experiment_runs(exp_id, limit=limit, offset=offset)
    return {"ok": True, "data": {"items": runs, "total": total, "limit": limit, "offset": offset}}


# ---------------------------------------------------------------------------
# GET /api/experiments/{exp_id}/export?format=json|csv
# ---------------------------------------------------------------------------

@router.get("/api/experiments/{exp_id}/export")
def api_export_experiment(
    exp_id: int,
    request: Request,
    fmt: str = Query(default="json", alias="format"),
):
    _require_scope(request, "experiments:read")

    exp = get_experiment(exp_id)
    if not exp:
        raise HTTPException(status_code=404, detail="Experiment not found")

    _check_experiment_access(request, exp)

    runs, _total = list_experiment_runs(exp_id, limit=10000, offset=0)

    if fmt == "csv":
        buf = io.StringIO()
        if runs:
            # Flatten scores into top-level columns
            all_score_keys: set[str] = set()
            for r in runs:
                if isinstance(r.get("scores"), dict):
                    all_score_keys.update(r["scores"].keys())
            score_cols = sorted(all_score_keys - {"parse_error"})

            base_fields = [
                "id", "experiment_id", "variant_label", "model_spec",
                "input_text", "output_text", "latency_ms", "status", "created_at",
            ]
            fieldnames = base_fields + [f"score_{k}" for k in score_cols]
            writer = csv.DictWriter(buf, fieldnames=fieldnames, extrasaction="ignore")
            writer.writeheader()
            for r in runs:
                row = {k: r.get(k, "") for k in base_fields}
                scores = r.get("scores") or {}
                for k in score_cols:
                    row[f"score_{k}"] = scores.get(k, "")
                writer.writerow(row)
        return StreamingResponse(
            iter([buf.getvalue()]),
            media_type="text/csv",
            headers={"Content-Disposition": f"attachment; filename=experiment_{exp_id}.csv"},
        )

    # Default: JSON
    payload = {"experiment": exp, "runs": runs}
    return StreamingResponse(
        iter([json.dumps(payload, ensure_ascii=False, indent=2)]),
        media_type="application/json",
        headers={"Content-Disposition": f"attachment; filename=experiment_{exp_id}.json"},
    )
