Create-knowledge-find-decode/acquisition/repositories/postgres.py

"""PostgreSQL implementation of the formal acquisition repository."""
from __future__ import annotations

from typing import Any
from uuid import UUID

import psycopg2.extras

from acquisition.domain import (
    AcquisitionJob,
    AcquisitionRun,
    CandidateItem,
    ItemClassification,
    MediaAsset,
    Query,
    QueryBatch,
)


Json = psycopg2.extras.Json
psycopg2.extras.register_uuid()


class PostgresAcquisitionRepository:
    """Repository backed by the formal cloud PostgreSQL schema.

    The repository does not commit by itself; callers own transaction scope via
    core.db_session.transaction or an equivalent connection boundary.
    """

    def __init__(self, conn: Any):
        self.conn = conn

    def _one(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any]:
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
            if row is None:
                raise RuntimeError("expected one row, got none")
            return dict(row)

    def _all(self, sql: str, params: tuple[Any, ...]) -> list[dict[str, Any]]:
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            return [dict(row) for row in cur.fetchall()]

    def create_query_batch(
        self,
        *,
        name: str,
        source_type: str = "manual",
        generation_method: str | None = None,
        target_platforms: list[str] | None = None,
        status: str = "draft",
        metadata: dict[str, Any] | None = None,
    ) -> QueryBatch:
        row = self._one(
            """
            INSERT INTO query_batches(
                name, source_type, generation_method, target_platforms,
                status, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                name,
                source_type,
                generation_method,
                target_platforms or [],
                status,
                Json(metadata or {}),
            ),
        )
        return QueryBatch.model_validate(row)

    def add_query(
        self,
        *,
        batch_id: UUID | None,
        query_text: str,
        axes: dict[str, Any] | None = None,
        keep: bool | None = None,
        filter_reason: str | None = None,
        status: str = "draft",
        sort_order: int = 0,
        metadata: dict[str, Any] | None = None,
    ) -> Query:
        row = self._one(
            """
            INSERT INTO queries(
                batch_id, query_text, axes, keep, filter_reason,
                status, sort_order, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                batch_id,
                query_text,
                Json(axes or {}),
                keep,
                filter_reason,
                status,
                sort_order,
                Json(metadata or {}),
            ),
        )
        return Query.model_validate(row)

    def list_queries_for_batch(
        self,
        batch_id: UUID,
        *,
        keep: bool | None = None,
    ) -> list[Query]:
        if keep is None:
            rows = self._all(
                """
                SELECT * FROM queries
                WHERE batch_id = %s
                ORDER BY sort_order, created_at
                """,
                (batch_id,),
            )
        else:
            rows = self._all(
                """
                SELECT * FROM queries
                WHERE batch_id = %s AND keep IS NOT DISTINCT FROM %s
                ORDER BY sort_order, created_at
                """,
                (batch_id, keep),
            )
        return [Query.model_validate(row) for row in rows]

    def get_query_batch(self, batch_id: UUID) -> QueryBatch:
        row = self._one("SELECT * FROM query_batches WHERE id = %s", (batch_id,))
        return QueryBatch.model_validate(row)

    def create_acquisition_run(
        self,
        *,
        batch_id: UUID | None = None,
        run_key: str | None = None,
        status: str = "pending",
        note: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AcquisitionRun:
        row = self._one(
            """
            INSERT INTO acquisition_runs(batch_id, run_key, status, note, metadata)
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (run_key) DO UPDATE SET
                status = EXCLUDED.status,
                note = EXCLUDED.note,
                metadata = EXCLUDED.metadata
            RETURNING *
            """,
            (batch_id, run_key, status, note, Json(metadata or {})),
        )
        return AcquisitionRun.model_validate(row)

    def ensure_acquisition_job(
        self,
        *,
        run_id: UUID,
        query_id: UUID | None,
        platform: str,
        search_limit: int | None = None,
        display_limit: int | None = None,
        status: str = "pending",
        metadata: dict[str, Any] | None = None,
    ) -> AcquisitionJob:
        row = self._one(
            """
            INSERT INTO acquisition_jobs(
                run_id, query_id, platform, search_limit,
                display_limit, status, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (run_id, query_id, platform) DO UPDATE SET
                search_limit = EXCLUDED.search_limit,
                display_limit = EXCLUDED.display_limit,
                metadata = EXCLUDED.metadata
            RETURNING *
            """,
            (
                run_id,
                query_id,
                platform,
                search_limit,
                display_limit,
                status,
                Json(metadata or {}),
            ),
        )
        return AcquisitionJob.model_validate(row)

    def update_acquisition_job(
        self,
        job_id: UUID,
        *,
        status: str,
        attempt_count: int | None = None,
        error_message: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AcquisitionJob:
        row = self._one(
            """
            UPDATE acquisition_jobs SET
                status = %s,
                attempt_count = COALESCE(%s, attempt_count),
                error_message = %s,
                metadata = CASE WHEN %s THEN %s ELSE metadata END
            WHERE id = %s
            RETURNING *
            """,
            (
                status,
                attempt_count,
                error_message,
                metadata is not None,
                Json(metadata or {}),
                job_id,
            ),
        )
        return AcquisitionJob.model_validate(row)

    def upsert_candidate_item(
        self,
        *,
        platform: str,
        job_id: UUID | None = None,
        query_id: UUID | None = None,
        platform_item_id: str | None = None,
        canonical_url: str | None = None,
        content_type: str | None = None,
        title: str | None = None,
        author_name: str | None = None,
        raw_summary: str | None = None,
        status: str = "candidate",
        source_payload: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        error_message: str | None = None,
    ) -> CandidateItem:
        existing_id = None
        if platform_item_id:
            row = self._one_or_none(
                """
                SELECT id FROM candidate_items
                WHERE platform = %s AND platform_item_id = %s
                ORDER BY created_at DESC
                LIMIT 1
                """,
                (platform, platform_item_id),
            )
            existing_id = row["id"] if row else None

        if existing_id:
            row = self._one(
                """
                UPDATE candidate_items SET
                    job_id = %s,
                    query_id = %s,
                    canonical_url = %s,
                    content_type = %s,
                    title = %s,
                    author_name = %s,
                    raw_summary = %s,
                    status = %s,
                    source_payload = %s,
                    metadata = %s,
                    error_message = %s
                WHERE id = %s
                RETURNING *
                """,
                (
                    job_id,
                    query_id,
                    canonical_url,
                    content_type,
                    title,
                    author_name,
                    raw_summary,
                    status,
                    Json(source_payload or {}),
                    Json(metadata or {}),
                    error_message,
                    existing_id,
                ),
            )
        else:
            row = self._one(
                """
                INSERT INTO candidate_items(
                    job_id, query_id, platform, platform_item_id, canonical_url,
                    content_type, title, author_name, raw_summary, status,
                    source_payload, metadata, error_message
                )
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                RETURNING *
                """,
                (
                    job_id,
                    query_id,
                    platform,
                    platform_item_id,
                    canonical_url,
                    content_type,
                    title,
                    author_name,
                    raw_summary,
                    status,
                    Json(source_payload or {}),
                    Json(metadata or {}),
                    error_message,
                ),
            )
        return CandidateItem.model_validate(row)

    def _one_or_none(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any] | None:
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
            return dict(row) if row else None

    def add_media_asset(
        self,
        *,
        item_id: UUID,
        media_type: str,
        source_url: str | None = None,
        oss_url: str | None = None,
        cdn_url: str | None = None,
        position: int = 0,
        status: str = "pending",
        source_payload: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> MediaAsset:
        row = self._one(
            """
            INSERT INTO media_assets(
                item_id, media_type, source_url, oss_url, cdn_url,
                position, status, source_payload, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                item_id,
                media_type,
                source_url,
                oss_url,
                cdn_url,
                position,
                status,
                Json(source_payload or {}),
                Json(metadata or {}),
            ),
        )
        return MediaAsset.model_validate(row)

    def add_item_classification(
        self,
        *,
        item_id: UUID,
        is_creation_knowledge: bool | None = None,
        label: str | None = None,
        confidence: float | None = None,
        reason: str | None = None,
        model_name: str | None = None,
        prompt_version: str | None = None,
        result_payload: dict[str, Any] | None = None,
        status: str = "pending",
        error_message: str | None = None,
    ) -> ItemClassification:
        row = self._one(
            """
            INSERT INTO item_classifications(
                item_id, is_creation_knowledge, label, confidence, reason,
                model_name, prompt_version, result_payload, status, error_message
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                item_id,
                is_creation_knowledge,
                label,
                confidence,
                reason,
                model_name,
                prompt_version,
                Json(result_payload or {}),
                status,
                error_message,
            ),
        )
        return ItemClassification.model_validate(row)

    def get_run_summary(self, run_id: UUID) -> dict[str, Any]:
        summary = self._one(
            """
            SELECT
                ar.id,
                ar.run_key,
                ar.batch_id,
                ar.status,
                ar.started_at,
                ar.finished_at,
                COUNT(DISTINCT q.id)::int AS query_count,
                COUNT(DISTINCT aj.id)::int AS job_count,
                COUNT(DISTINCT ci.id)::int AS candidate_count,
                COUNT(DISTINCT ic.id) FILTER (
                    WHERE ic.is_creation_knowledge IS TRUE
                )::int AS creation_hit_count
            FROM acquisition_runs ar
            LEFT JOIN acquisition_jobs aj ON aj.run_id = ar.id
            LEFT JOIN queries q ON q.id = aj.query_id
            LEFT JOIN candidate_items ci ON ci.job_id = aj.id
            LEFT JOIN item_classifications ic ON ic.item_id = ci.id
            WHERE ar.id = %s
            GROUP BY ar.id
            """,
            (run_id,),
        )
        summary["queries"] = self._all(
            """
            SELECT
                q.id AS query_id,
                q.query_text,
                COUNT(DISTINCT aj.id)::int AS job_count,
                COUNT(DISTINCT ci.id)::int AS candidate_count,
                COUNT(DISTINCT ic.id) FILTER (
                    WHERE ic.is_creation_knowledge IS TRUE
                )::int AS creation_hit_count,
                jsonb_object_agg(
                    aj.platform,
                    jsonb_build_object(
                        'status', aj.status,
                        'attempt_count', aj.attempt_count,
                        'display_limit', aj.display_limit,
                        'search_limit', aj.search_limit,
                        'error_message', aj.error_message
                    )
                ) FILTER (WHERE aj.id IS NOT NULL) AS platforms
            FROM queries q
            JOIN acquisition_jobs aj ON aj.query_id = q.id
            LEFT JOIN candidate_items ci ON ci.job_id = aj.id
            LEFT JOIN item_classifications ic ON ic.item_id = ci.id
            WHERE aj.run_id = %s
            GROUP BY q.id, q.query_text, q.sort_order
            ORDER BY q.sort_order, q.query_text
            """,
            (run_id,),
        )
        return summary

    def get_query_detail(self, *, run_id: UUID, query_id: UUID) -> dict[str, Any]:
        query = self._one("SELECT * FROM queries WHERE id = %s", (query_id,))
        jobs = self._all(
            """
            SELECT * FROM acquisition_jobs
            WHERE run_id = %s AND query_id = %s
            ORDER BY platform
            """,
            (run_id, query_id),
        )
        items = self._all(
            """
            SELECT ci.* FROM candidate_items ci
            JOIN acquisition_jobs aj ON aj.id = ci.job_id
            WHERE aj.run_id = %s AND ci.query_id = %s
            ORDER BY ci.platform, ci.created_at
            """,
            (run_id, query_id),
        )
        item_ids = [row["id"] for row in items]
        media: list[dict[str, Any]] = []
        classifications: list[dict[str, Any]] = []
        if item_ids:
            media = self._all(
                """
                SELECT * FROM media_assets
                WHERE item_id = ANY(%s)
                ORDER BY item_id, position
                """,
                (item_ids,),
            )
            classifications = self._all(
                """
                SELECT * FROM item_classifications
                WHERE item_id = ANY(%s)
                ORDER BY created_at DESC
                """,
                (item_ids,),
            )
        return {
            "query": query,
            "jobs": jobs,
            "items": items,
            "media_assets": media,
            "classifications": classifications,
        }

    def list_creation_candidate_items(
        self,
        *,
        run_id: UUID | None = None,
        limit: int = 100,
    ) -> list[CandidateItem]:
        if run_id is None:
            rows = self._all(
                """
                SELECT ci.* FROM candidate_items ci
                JOIN item_classifications ic ON ic.item_id = ci.id
                WHERE ic.is_creation_knowledge IS TRUE
                ORDER BY ci.created_at
                LIMIT %s
                """,
                (limit,),
            )
        else:
            rows = self._all(
                """
                SELECT ci.* FROM candidate_items ci
                JOIN acquisition_jobs aj ON aj.id = ci.job_id
                JOIN item_classifications ic ON ic.item_id = ci.id
                WHERE aj.run_id = %s AND ic.is_creation_knowledge IS TRUE
                ORDER BY ci.created_at
                LIMIT %s
                """,
                (run_id, limit),
            )
        return [CandidateItem.model_validate(row) for row in rows]

    def get_candidate_item(self, item_id: UUID) -> CandidateItem:
        row = self._one("SELECT * FROM candidate_items WHERE id = %s", (item_id,))
        return CandidateItem.model_validate(row)

    def list_media_assets_for_item(self, item_id: UUID) -> list[MediaAsset]:
        rows = self._all(
            """
            SELECT * FROM media_assets
            WHERE item_id = %s
            ORDER BY position
            """,
            (item_id,),
        )
        return [MediaAsset.model_validate(row) for row in rows]