Blame for annotation_coverage.py

"""Public-safe associated video coverage for eligible representative batches.

Every active member contributes its full measured duration to the bucket of its
collection start. These durations describe associated video coverage, without
asserting that a representative frame's labels persist throughout the videos.
"""
from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from datetime import date
from decimal import Decimal
import math
import re

import psycopg

from annotation_statistics import validate_public_labels
from live_inventory import _period_for_hour, _valid_number


_ELIGIBLE_BATCH_SQL = """
SELECT b.batch_id, b.labels, b.schema_version
FROM aerial.annotation_batches b
JOIN aerial.videos representative
  ON representative.sha256 = b.representative_video_sha256
WHERE b.processing_info->>'status' = 'completed'
  AND b.validation_passed IS TRUE
  AND b.review_status IS DISTINCT FROM 'rejected'
  AND representative.deleted_at IS NULL
  AND EXISTS (
      SELECT 1 FROM aerial.annotation_batch_videos bv
      WHERE bv.batch_id = b.batch_id
        AND bv.video_sha256 = representative.sha256
  )
ORDER BY b.batch_id
"""
_MEMBER_SQL = """
SELECT bv.batch_id, v.sha256,
       (v.collected_at AT TIME ZONE 'Asia/Shanghai')::date,
       EXTRACT(HOUR FROM v.collected_at AT TIME ZONE 'Asia/Shanghai')::integer,
       v.codec, v.duration_seconds
FROM aerial.annotation_batch_videos bv
JOIN aerial.videos v ON v.sha256 = bv.video_sha256
WHERE bv.batch_id = ANY(%s)
  AND v.deleted_at IS NULL
ORDER BY bv.batch_id, v.sha256
"""
_PUBLIC_CODE = re.compile(r'B-[0-9a-f]{16}')


@dataclass(frozen=True)
class CoverageSlice:
    collected_on: date | None
    period: str
    codec: str
    known_seconds: float
    known_videos: int
    unknown_videos: int


@dataclass(frozen=True)
class BatchCoverage:
    batch_code: str
    slices: tuple[CoverageSlice, ...]


@dataclass(frozen=True)
class CoverageCatalog:
    records: tuple[BatchCoverage, ...]


def _slice_key(item: CoverageSlice) -> tuple[bool, date, str, str]:
    return item.collected_on is None, item.collected_on or date.max, item.period, item.codec


def _validated_code(value: object) -> str:
    if not isinstance(value, str) or not _PUBLIC_CODE.fullmatch(value):
        raise ValueError('Invalid coverage batch code')
    return value


def align_coverage_catalog(catalog: CoverageCatalog, batch_codes: Iterable[str]) -> CoverageCatalog:
    """Match independent source coverage to the exact exported batch population.

    Empty slices mark unavailable member coverage; a known zero remains a slice
    with a positive known-video count. Reject duplicate source or export codes
    rather than silently hiding an inconsistent publication population.
    """
    records = {}
    for record in catalog.records:
        code = _validated_code(record.batch_code)
        if code in records:
            raise ValueError('Duplicate coverage batch code')
        records[code] = record
    selected = set()
    for value in batch_codes:
        code = _validated_code(value)
        if code in selected:
            raise ValueError('Duplicate coverage batch code')
        selected.add(code)
    return CoverageCatalog(tuple(
        BatchCoverage(code, tuple(sorted(records[code].slices, key=_slice_key)))
        if code in records else BatchCoverage(code, ())
        for code in sorted(selected)
    ))


def fetch_coverage_catalog(dsn: str) -> CoverageCatalog:
    """Read eligible labels once per batch, then only active member metadata.

    Private identifiers are used locally for membership checks and discarded.
    Duplicate membership rows inside a batch are counted once. Conflicting
    active memberships across eligible batches fail with an identifier-free
    error. Missing, negative and nonfinite durations count as unknown videos;
    observed zero counts as a known duration. No duration is inferred.
    """
    # batch_updates imports the label validator too; keep a single safe B-code
    # contract without creating an import-time dependency cycle.
    from batch_updates import _public_code

    with psycopg.connect(dsn, options='-c default_transaction_read_only=on') as connection:
        with connection.cursor() as cursor:
            cursor.execute(_ELIGIBLE_BATCH_SQL)
            eligible = []
            for batch_id, raw_labels, schema_version in cursor.fetchall():
                try:
                    if schema_version not in ('aerial_annotation_v1', 'aerial_annotation_v2'):
                        raise ValueError('Unsupported source annotation schema')
                    validate_public_labels(raw_labels, schema_version)
                except ValueError:
                    continue
                eligible.append(batch_id)
            if not eligible:
                return CoverageCatalog(())
            cursor.execute(_MEMBER_SQL, (eligible,))
            rows = cursor.fetchall()

    buckets = {batch_id: {} for batch_id in eligible}
    ownership = {}
    for batch_id, video_id, collected_on, hour, codec, seconds in rows:
        if video_id in ownership:
            if ownership[video_id] != batch_id:
                raise ValueError('Conflicting active video batch membership')
            continue
        ownership[video_id] = batch_id
        key = (collected_on, _period_for_hour(hour),
               codec.strip() if isinstance(codec, str) and codec.strip() else '未知')
        bucket = buckets[batch_id].setdefault(key, [Decimal(0), 0, 0])
        duration = _valid_number(seconds)
        if duration is None:
            bucket[2] += 1
        else:
            bucket[0] += duration
            bucket[1] += 1

    records = []
    for batch_id, batch_buckets in buckets.items():
        slices = []
        for (collected_on, period, codec), (seconds, known, unknown) in batch_buckets.items():
            total = float(seconds)
            if not math.isfinite(total):
                raise ValueError('Invalid coverage duration total')
            slices.append(CoverageSlice(collected_on, period, codec, total, known, unknown))
        records.append(BatchCoverage(_public_code(batch_id), tuple(sorted(slices, key=_slice_key))))
    return CoverageCatalog(tuple(sorted(records, key=lambda record: record.batch_code)))