Create-knowledge-find-decode/app/routes/acquisition.py

"""Formal acquisition API routes."""
from __future__ import annotations

from collections import defaultdict
from typing import Any
from uuid import UUID

from fastapi import APIRouter, Depends

from acquisition.repositories.base import AcquisitionRepository
from app.dependencies import get_acquisition_repository
from app.schemas import (
    AcquisitionRunSummarySchema,
    CandidateItemSchema,
    ItemClassificationSchema,
    MediaAssetSchema,
    QueryBatchSchema,
    QuerySchema,
)

router = APIRouter(prefix="/api", tags=["acquisition"])


def _model_dump(row: Any) -> dict[str, Any]:
    if isinstance(row, dict):
        return dict(row)
    if hasattr(row, "model_dump"):
        return row.model_dump()
    return dict(row)


def _normalize_query_detail(detail: dict[str, Any]) -> dict[str, Any]:
    media_by_item: dict[Any, list[dict[str, Any]]] = defaultdict(list)
    for media in detail.get("media_assets") or []:
        media_by_item[media["item_id"]].append(media)

    classification_by_item: dict[Any, dict[str, Any]] = {}
    for row in detail.get("classifications") or []:
        classification_by_item.setdefault(row["item_id"], row)

    decode_by_item: dict[Any, dict[str, Any]] = {}
    for row in detail.get("decode_summaries") or []:
        decode_by_item.setdefault(row["item_id"], row)

    items: list[dict[str, Any]] = []
    platforms: dict[str, dict[str, Any]] = {}
    for raw_job in detail.get("jobs") or []:
        job = _model_dump(raw_job)
        platform = job.get("platform")
        if not platform:
            continue
        platforms.setdefault(
            platform,
            {
                "platform": platform,
                "status": job.get("status") or "pending",
                "attempt_count": job.get("attempt_count"),
                "display_limit": job.get("display_limit"),
                "search_limit": job.get("search_limit"),
                "error_message": job.get("error_message"),
                "items": [],
            },
        )
    for row in detail.get("items") or []:
        item = _model_dump(row)
        item["media_assets"] = [
            MediaAssetSchema.model_validate(media).model_dump(mode="json")
            for media in media_by_item.get(item["id"], [])
        ]
        cls = classification_by_item.get(item["id"])
        item["classification"] = (
            ItemClassificationSchema.model_validate(cls).model_dump(mode="json")
            if cls
            else None
        )
        item["decode_summary"] = _model_dump(decode_by_item[item["id"]]) if item["id"] in decode_by_item else None
        payload = CandidateItemSchema.model_validate(item).model_dump(mode="json")
        items.append(payload)
        group = platforms.setdefault(
            payload["platform"],
            {"platform": payload["platform"], "status": "done", "items": []},
        )
        if group.get("status") in {None, "pending"}:
            group["status"] = "done"
        group["items"].append(payload)
    return {
        "query": QuerySchema.model_validate(detail["query"]).model_dump(mode="json"),
        "jobs": [_model_dump(job) for job in detail.get("jobs") or []],
        "items": items,
        "platforms": platforms,
    }


@router.get("/query-batches/{batch_id}")
def query_batch(
    batch_id: UUID,
    repo: AcquisitionRepository = Depends(get_acquisition_repository),
) -> dict[str, Any]:
    batch = repo.get_query_batch(batch_id)
    queries = repo.list_queries_for_batch(batch_id)
    return {
        "batch": QueryBatchSchema.model_validate(batch).model_dump(mode="json"),
        "queries": [QuerySchema.model_validate(query).model_dump(mode="json") for query in queries],
    }


@router.get(
    "/acquisition/runs/{run_id}/summary",
    response_model=AcquisitionRunSummarySchema,
)
def acquisition_run_summary(
    run_id: UUID,
    repo: AcquisitionRepository = Depends(get_acquisition_repository),
) -> AcquisitionRunSummarySchema:
    return AcquisitionRunSummarySchema.model_validate(repo.get_run_summary(run_id))


@router.get("/acquisition/runs/{run_id}/queries/{query_id}")
def acquisition_query_detail(
    run_id: UUID,
    query_id: UUID,
    repo: AcquisitionRepository = Depends(get_acquisition_repository),
) -> dict[str, Any]:
    return _normalize_query_detail(repo.get_query_detail(run_id=run_id, query_id=query_id))