"""
Experiment Log API — AI version experiment tracking.
"""
import logging
import uuid
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from typing import List, Optional

from fastapi import APIRouter, HTTPException
from psycopg_pool import AsyncConnectionPool
from pydantic import BaseModel

from config import CHECKPOINT_POSTGRES_URL

VN_TZ = timezone(timedelta(hours=7))
_pool: AsyncConnectionPool | None = None


async def _get_pool() -> AsyncConnectionPool:
    global _pool
    if _pool is None:
        _pool = AsyncConnectionPool(CHECKPOINT_POSTGRES_URL, min_size=1, max_size=5, open=False)
        await _pool.open()
    return _pool


def _now() -> datetime:
    return datetime.now(VN_TZ)

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/dashboard", tags=["Experiment Log"])


async def _query(q: str, params: tuple = ()) -> list[dict]:
    pool = await _get_pool()
    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            await cur.execute(q, params)
            cols = [d[0] for d in cur.description] if cur.description else []
            rows = await cur.fetchall()
            return [dict(zip(cols, row)) for row in rows]


async def _execute(q: str, params: tuple = ()):
    pool = await _get_pool()
    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            await cur.execute(q, params)
        await conn.commit()


def _serialize(row: dict) -> dict:
    r = {**row}
    for k in ("created_at", "updated_at"):
        if k in r and r[k]:
            r[k] = r[k].isoformat()
    # Convert Decimal to float
    for k in ("score", "latency_avg", "cost_per_req", "error_rate"):
        if k in r and r[k] is not None:
            r[k] = float(r[k])
    return r


# ── Models ──

class ExperimentCreate(BaseModel):
    version: str
    title: str
    content: str = ""
    tags: List[str] = []
    score: Optional[float] = None
    latency_avg: Optional[float] = None
    cost_per_req: Optional[float] = None
    error_rate: Optional[float] = None
    model_name: str = ""
    status: str = "draft"
    pinned: bool = False


class ExperimentUpdate(BaseModel):
    version: Optional[str] = None
    title: Optional[str] = None
    content: Optional[str] = None
    tags: Optional[List[str]] = None
    score: Optional[float] = None
    latency_avg: Optional[float] = None
    cost_per_req: Optional[float] = None
    error_rate: Optional[float] = None
    model_name: Optional[str] = None
    status: Optional[str] = None
    pinned: Optional[bool] = None


# ── Endpoints ──

@router.get("/experiments")
async def list_experiments(tag: Optional[str] = None, status: Optional[str] = None):
    """List all experiments. Optional filter by tag or status."""
    q = "SELECT * FROM experiment_log WHERE 1=1"
    params = []
    if tag:
        q += " AND %s = ANY(tags)"
        params.append(tag)
    if status:
        q += " AND status = %s"
        params.append(status)
    q += " ORDER BY pinned DESC, created_at DESC"
    rows = await _query(q, tuple(params))
    items = [_serialize(r) for r in rows]

    # Collect all tags with counts
    tag_counts = defaultdict(int)
    all_rows = await _query("SELECT tags FROM experiment_log")
    for r in all_rows:
        for t in (r.get("tags") or []):
            tag_counts[t] += 1

    # Activity heatmap: date → count
    activity_rows = await _query(
        "SELECT DATE(created_at AT TIME ZONE 'Asia/Ho_Chi_Minh') as d, COUNT(*) as c FROM experiment_log GROUP BY d"
    )
    activity = {str(r["d"]): int(r["c"]) for r in activity_rows}

    return {
        "status": "success",
        "experiments": items,
        "total": len(items),
        "tags": dict(sorted(tag_counts.items(), key=lambda x: -x[1])),
        "activity": activity,
    }


@router.post("/experiments")
async def create_experiment(item: ExperimentCreate):
    """Create a new experiment entry."""
    eid = str(uuid.uuid4())[:8]
    now = _now()
    await _execute(
        """INSERT INTO experiment_log
           (id, version, title, content, tags, score, latency_avg, cost_per_req, error_rate, model_name, status, pinned, author, created_at, updated_at)
           VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,'Admin',%s,%s)""",
        (eid, item.version, item.title, item.content, item.tags,
         item.score, item.latency_avg, item.cost_per_req, item.error_rate,
         item.model_name, item.status, item.pinned, now, now),
    )
    return {"status": "success", "id": eid}


@router.put("/experiments/{exp_id}")
async def update_experiment(exp_id: str, update: ExperimentUpdate):
    """Update an experiment."""
    sets, vals = [], []
    for field in ("version", "title", "content", "score", "latency_avg",
                  "cost_per_req", "error_rate", "model_name", "status", "pinned"):
        val = getattr(update, field, None)
        if val is not None:
            sets.append(f"{field} = %s")
            vals.append(val)
    if update.tags is not None:
        sets.append("tags = %s")
        vals.append(update.tags)
    if not sets:
        raise HTTPException(status_code=400, detail="Nothing to update")
    sets.append("updated_at = %s")
    vals.append(_now())
    vals.append(exp_id)
    await _execute(f"UPDATE experiment_log SET {', '.join(sets)} WHERE id = %s", tuple(vals))
    return {"status": "success"}


@router.delete("/experiments/{exp_id}")
async def delete_experiment(exp_id: str):
    """Delete an experiment."""
    await _execute("DELETE FROM experiment_log WHERE id = %s", (exp_id,))
    return {"status": "success"}


@router.patch("/experiments/{exp_id}/pin")
async def toggle_experiment_pin(exp_id: str):
    """Toggle pin status."""
    rows = await _query("SELECT pinned FROM experiment_log WHERE id = %s", (exp_id,))
    if not rows:
        raise HTTPException(status_code=404, detail="Not found")
    new_pin = not rows[0]["pinned"]
    await _execute("UPDATE experiment_log SET pinned = %s, updated_at = %s WHERE id = %s",
                   (new_pin, _now(), exp_id))
    return {"status": "success", "pinned": new_pin}


@router.get("/experiments/compare")
async def compare_experiments(id1: str, id2: str):
    """Compare two experiments side-by-side."""
    rows = await _query("SELECT * FROM experiment_log WHERE id IN (%s, %s)", (id1, id2))
    if len(rows) < 2:
        raise HTTPException(status_code=404, detail="One or both experiments not found")
    items = [_serialize(r) for r in rows]
    return {"status": "success", "experiments": items}
