reTraining/backend/api/counting_bench.py

"""Counting bench routes: the scored table and the background recount (REQ-150…153)."""

from typing import List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from backend import archive_index, counting_bench, jobs
from backend.api.common import project_or_404

router = APIRouter(tags=["counting-bench"])


class GroundTruthRequest(BaseModel):
    video_rel: str
    # Null clears it, which has to be distinguishable from "not sent" — a
    # cleared row must drop back out of the accuracy totals.
    ground_truth: Optional[int] = None


class ClockRequest(BaseModel):
    video_rel: str
    # ISO "2026-08-13 08:27:28"; null clears a manual entry and lets a rescan
    # read the overlay again.
    started_at: Optional[str] = None


class ScanRequest(BaseModel):
    rescan: bool = False


class CountRequest(BaseModel):
    video_rels: List[str] = []
    model_path: str = ""
    model_version_id: Optional[int] = None
    params: Optional[dict] = None
    recount: bool = False
    all_videos: bool = False


@router.get("/api/projects/{project_id}/counting-bench")
def listing(project_id: int, date: Optional[str] = None) -> dict:
    project_or_404(project_id)
    try:
        payload = counting_bench.listing(project_id, date)
    except counting_bench.CountingBenchError as exc:
        raise HTTPException(400, str(exc))
    payload["defaults"] = counting_bench.DEFAULTS
    running = [job for job in (j.to_dict() for j in jobs.listing(project_id))
               if job["status"] in ("queued", "running")]
    payload["active_job"] = next((j for j in running if j["type"] == "count"), None)
    payload["scan_job"] = next((j for j in running if j["type"] == "clock-scan"), None)
    return payload


class TruckScanRequest(BaseModel):
    model_path: str = ""
    model_version_id: Optional[int] = None
    rescan: bool = False


@router.post("/api/projects/{project_id}/archive/truck-scan")
def truck_scan(project_id: int, body: TruckScanRequest) -> dict:
    """Check every recording actually contains a truck (REQ-166)."""
    project_or_404(project_id)
    from backend.api.live_count import available_models

    path = body.model_path
    models = available_models(project_id)["models"]
    if not path and body.model_version_id is not None:
        path = next((m["path"] for m in models
                     if m["version_id"] == body.model_version_id), "")
    # No model named: the newest trained version, which is what "v4" means.
    if not path:
        path = models[0]["path"] if models else ""
    if not path:
        raise HTTPException(400, "This project has no model to detect trucks with")
    try:
        return archive_index.queue_truck_scan(project_id, path, rescan=body.rescan)
    except archive_index.ArchiveIndexError as exc:
        raise HTTPException(400, str(exc))


@router.get("/api/projects/{project_id}/archive/cycles")
def archive_cycles(project_id: int) -> dict:
    """The archive grouped into 06:00-to-05:59 cycles (REQ-165)."""
    project_or_404(project_id)
    try:
        return {"cycles": archive_index.cycles(project_id)}
    except archive_index.ArchiveIndexError as exc:
        raise HTTPException(400, str(exc))


@router.get("/api/projects/{project_id}/archive/cycles/{cycle}")
def archive_cycle_videos(project_id: int, cycle: str) -> dict:
    project_or_404(project_id)
    try:
        return {"videos": archive_index.cycle_videos(project_id, cycle)}
    except archive_index.ArchiveIndexError as exc:
        raise HTTPException(400, str(exc))


@router.post("/api/projects/{project_id}/counting-bench/scan-clock")
def scan_clock(project_id: int, body: ScanRequest) -> dict:
    """Read the burned-in timestamp off every recording (REQ-160)."""
    project_or_404(project_id)
    try:
        return archive_index.queue_scan(project_id, rescan=body.rescan)
    except archive_index.ArchiveIndexError as exc:
        raise HTTPException(400, str(exc))


@router.patch("/api/projects/{project_id}/counting-bench/clock")
def set_clock(project_id: int, body: ClockRequest) -> dict:
    """Hand-correct a start time the reader could not get right."""
    import datetime

    project_or_404(project_id)
    stamp = None
    if body.started_at:
        for shape in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M"):
            try:
                stamp = datetime.datetime.strptime(body.started_at, shape)
                break
            except ValueError:
                continue
        if stamp is None:
            raise HTTPException(400, "Use the form 2026-08-13 08:27:28")
    return archive_index.set_manual(
        project_id, body.video_rel,
        stamp.strftime("%Y-%m-%d %H:%M:%S") if stamp else None)


@router.patch("/api/projects/{project_id}/counting-bench/ground-truth")
def set_ground_truth(project_id: int, body: GroundTruthRequest) -> dict:
    project_or_404(project_id)
    if body.ground_truth is not None and body.ground_truth < 0:
        raise HTTPException(400, "A ground truth count cannot be negative")
    return counting_bench.set_ground_truth(project_id, body.video_rel, body.ground_truth)


@router.post("/api/projects/{project_id}/counting-bench/run")
def run(project_id: int, body: CountRequest) -> dict:
    project_or_404(project_id)

    path = body.model_path
    if not path and body.model_version_id is not None:
        from backend.api.live_count import available_models
        for item in available_models(project_id)["models"]:
            if item["version_id"] == body.model_version_id:
                path = item["path"]
                break
    if not path:
        raise HTTPException(400, "Pick a model to count with")

    rels = body.video_rels
    if body.all_videos:
        rels = [row["video_rel"] for row in counting_bench.listing(project_id)["rows"]]

    try:
        return counting_bench.queue(project_id, rels, path,
                                    params=body.params, recount=body.recount)
    except counting_bench.CountingBenchError as exc:
        raise HTTPException(400, str(exc))