"""Heads API (#114): train + inspect the per-concept heads that power
suggestions (replacing Camie + centroid).

POST /api/heads/train   — (re)train all eligible heads (one run at a time).
GET  /api/heads         — status: head count, last-trained, running run, the
                          per-concept head table (strength + auto-apply ready),
                          and recent training runs. The card rehydrates from
                          here so status survives navigation.
"""

from quart import Blueprint, jsonify, request
from sqlalchemy import desc, func, select

from ..extensions import get_session
from ..models import (
    HeadAutoApplyRun,
    HeadMetric,
    HeadMetricsSnapshot,
    HeadTrainingRun,
    Tag,
    TagHead,
)
from ..models.tag import image_tag
from ..services.ml.heads import (
    HeadAutoApplyAlreadyRunning,
    HeadAutoApplyDisabled,
    HeadTrainingAlreadyRunning,
    start_head_auto_apply_run,
    start_head_training_run,
)

heads_bp = Blueprint("heads", __name__, url_prefix="/api/heads")


def _serialize_run(run: HeadTrainingRun) -> dict:
    return {
        "id": run.id,
        "params": run.params,
        "status": run.status,
        "started_at": run.started_at.isoformat() if run.started_at else None,
        "finished_at": run.finished_at.isoformat() if run.finished_at else None,
        "n_trained": run.n_trained,
        "n_skipped": run.n_skipped,
        "error": run.error,
    }


@heads_bp.route("/train", methods=["POST"])
async def train():
    body = await request.get_json(silent=True) or {}
    params = body.get("params") or body or {}
    async with get_session() as session:
        try:
            run_id = await session.run_sync(
                lambda s: start_head_training_run(s, params)
            )
        except HeadTrainingAlreadyRunning as running:
            return jsonify({
                "error": "training_already_running",
                "running_id": int(running.args[0]),
            }), 409
        await session.commit()
    return jsonify({"run_id": run_id, "status": "running"}), 202


@heads_bp.route("", methods=["GET"])
async def status():
    async with get_session() as session:
        count, last_trained = (
            await session.execute(
                select(func.count(), func.max(TagHead.trained_at))
            )
        ).one()
        graduated = (
            await session.execute(
                select(func.count()).where(
                    TagHead.auto_apply_threshold.is_not(None)
                )
            )
        ).scalar_one()
        running = (
            await session.execute(
                select(HeadTrainingRun.id)
                .where(HeadTrainingRun.status == "running")
                .order_by(HeadTrainingRun.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        runs = (
            await session.execute(
                select(HeadTrainingRun)
                .order_by(HeadTrainingRun.id.desc())
                .limit(10)
            )
        ).scalars().all()
        # The per-concept table: strongest first, capped for the admin card.
        head_rows = (
            await session.execute(
                select(
                    TagHead.tag_id, Tag.name, Tag.kind,
                    TagHead.n_pos, TagHead.n_neg, TagHead.ap,
                    TagHead.precision_cv, TagHead.recall,
                    TagHead.auto_apply_threshold, TagHead.trained_at,
                )
                .join(Tag, Tag.id == TagHead.tag_id)
                .order_by(desc(TagHead.ap))
                .limit(500)
            )
        ).all()
        heads = [
            {
                "tag_id": r.tag_id,
                "name": r.name,
                "category": r.kind.value if hasattr(r.kind, "value") else str(r.kind),
                "n_pos": r.n_pos,
                "n_neg": r.n_neg,
                "ap": r.ap,
                "precision": r.precision_cv,
                "recall": r.recall,
                "auto_apply": r.auto_apply_threshold is not None,
                "trained_at": r.trained_at.isoformat() if r.trained_at else None,
            }
            for r in head_rows
        ]
    return jsonify({
        "head_count": count,
        "graduated_count": graduated,
        "last_trained_at": last_trained.isoformat() if last_trained else None,
        "running_id": running,
        "runs": [_serialize_run(r) for r in runs],
        "heads": heads,
    })


def _serialize_apply_run(run: HeadAutoApplyRun) -> dict:
    return {
        "id": run.id,
        "dry_run": run.dry_run,
        "status": run.status,
        "started_at": run.started_at.isoformat() if run.started_at else None,
        "finished_at": run.finished_at.isoformat() if run.finished_at else None,
        "n_applied": run.n_applied,
        "report": run.report,
        "error": run.error,
    }


@heads_bp.route("/auto-apply", methods=["POST"])
async def auto_apply():
    """Trigger an earned-auto-apply sweep. {dry_run:true} previews (writes
    nothing); a real sweep needs head_auto_apply_enabled on."""
    body = await request.get_json(silent=True) or {}
    params = {"dry_run": bool(body.get("dry_run", False))}
    async with get_session() as session:
        try:
            run_id = await session.run_sync(
                lambda s: start_head_auto_apply_run(s, params)
            )
        except HeadAutoApplyAlreadyRunning as running:
            return jsonify({
                "error": "auto_apply_already_running",
                "running_id": int(running.args[0]),
            }), 409
        except HeadAutoApplyDisabled:
            return jsonify({"error": "auto_apply_disabled"}), 400
        await session.commit()
    return jsonify({"run_id": run_id, "status": "running"}), 202


@heads_bp.route("/auto-apply", methods=["GET"])
async def auto_apply_status():
    async with get_session() as session:
        running = (
            await session.execute(
                select(HeadAutoApplyRun.id)
                .where(HeadAutoApplyRun.status == "running")
                .order_by(HeadAutoApplyRun.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        runs = (
            await session.execute(
                select(HeadAutoApplyRun)
                .order_by(HeadAutoApplyRun.id.desc())
                .limit(10)
            )
        ).scalars().all()
    return jsonify({
        "running_id": running,
        "runs": [_serialize_apply_run(r) for r in runs],
    })


@heads_bp.route("/metrics", methods=["GET"])
async def metrics():
    """Auto-apply observability: per-concept current counts (volume, misfires,
    under-fires, realized misfire rate, head quality) + the daily time-series so
    the operator can tune the precision target + support floor from real data."""
    async with get_session() as session:
        head_rows = (
            await session.execute(
                select(
                    TagHead.tag_id, Tag.name, TagHead.ap, TagHead.precision_cv,
                    TagHead.recall, TagHead.auto_apply_threshold, TagHead.n_pos,
                ).join(Tag, Tag.id == TagHead.tag_id)
            )
        ).all()
        heads = {r.tag_id: r for r in head_rows}
        metric_rows = (
            await session.execute(
                select(
                    HeadMetric.tag_id, HeadMetric.n_misfires, HeadMetric.n_underfires
                )
            )
        ).all()
        mets = {r.tag_id: r for r in metric_rows}
        applied = dict(
            (
                await session.execute(
                    select(image_tag.c.tag_id, func.count())
                    .where(image_tag.c.source == "head_auto")
                    .group_by(image_tag.c.tag_id)
                )
            ).all()
        )
        names = {r.tag_id: r.name for r in head_rows}
        # Names for metric-only tags (head pruned but corrections recorded).
        missing = [t for t in mets if t not in names]
        if missing:
            for tid, nm in (
                await session.execute(
                    select(Tag.id, Tag.name).where(Tag.id.in_(missing))
                )
            ).all():
                names[tid] = nm

        concepts = []
        for tid in set(heads) | set(mets):
            h = heads.get(tid)
            m = mets.get(tid)
            n_applied = applied.get(tid, 0)
            n_mis = m.n_misfires if m else 0
            denom = n_applied + n_mis
            concepts.append({
                "tag_id": tid,
                "name": names.get(tid, str(tid)),
                "n_auto_applied": n_applied,
                "n_misfires": n_mis,
                "n_underfires": m.n_underfires if m else 0,
                # Of everything this head ever auto-applied, the fraction you
                # removed — the misfire rate (null until something fired).
                "misfire_rate": round(n_mis / denom, 4) if denom else None,
                "ap": h.ap if h else None,
                "precision_cv": h.precision_cv if h else None,
                "recall": h.recall if h else None,
                "auto_apply": bool(h and h.auto_apply_threshold is not None),
                "n_pos": h.n_pos if h else None,
            })
        concepts.sort(key=lambda c: (c["n_misfires"], c["n_auto_applied"]), reverse=True)

        snaps = (
            await session.execute(
                select(HeadMetricsSnapshot)
                .order_by(HeadMetricsSnapshot.snapshot_at.desc())
                .limit(1000)
            )
        ).scalars().all()
    return jsonify({
        "concepts": concepts,
        "snapshots": [
            {
                "tag_id": s.tag_id,
                "name": s.name,
                "snapshot_at": s.snapshot_at.isoformat() if s.snapshot_at else None,
                "n_auto_applied": s.n_auto_applied,
                "n_misfires": s.n_misfires,
                "n_underfires": s.n_underfires,
                "ap": s.ap,
                "precision_cv": s.precision_cv,
                "recall": s.recall,
                "n_pos": s.n_pos,
            }
            for s in snaps
        ],
    })