Blame for annotation_progress.py

"""Read-only, public-safe progress for candidate aerial batch annotations.

This section is independent of official video statistics and never publishes
frames or writes review decisions. SQL projects an explicit public whitelist;
the returned immutable model validates those fields again before presentation.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
import json
import os
from pathlib import Path
from typing import Callable, Mapping
from uuid import UUID
from zoneinfo import ZoneInfo

import pandas as pd
import psycopg
import streamlit as st


_STATUSES = ("queued", "extracting", "labeling", "completed", "validation_failed", "failed", "skipped")
_STATES = frozenset(("observed", "not_observed", "unknown", "unlabeled", "not_applicable"))
_SCOPES = frozenset(("frame",))
_TIMEZONE = ZoneInfo("Asia/Shanghai")
_DATA = Path(__file__).resolve().parent / "data"
_ERRORS = {
    "source_unavailable": ("源视频暂不可读取", "Source video unavailable"),
    "no_valid_duration": ("缺少有效视频时长", "No valid video duration"),
    "frame_extract_failed": ("代表帧提取失败", "Representative frame extraction failed"),
    "api_auth": ("模型服务认证失败", "Model service authentication failed"),
    "api_quota": ("模型服务额度不足", "Model service quota exhausted"),
    "api_rate_limit": ("模型服务请求受限", "Model service rate limited"),
    "api_unavailable": ("模型服务暂不可用", "Model service unavailable"),
    "api_timeout": ("模型服务响应超时", "Model service timed out"),
    "invalid_json": ("模型返回格式未通过校验", "Model response format failed validation"),
    "invalid_labels": ("候选标签未通过校验", "Candidate labels failed validation"),
    "model_mismatch": ("返回模型与配置不一致", "Returned model differs from configuration"),
    "worker_exception": ("标注任务异常", "Annotation worker error"),
    "interrupted": ("标注任务已中断", "Annotation interrupted"),
    "source_group_changed": ("标注证据已更改，保留人工修改", "Annotation evidence changed; manual edits preserved"),
}

# A batch is eligible only while at least one associated video remains active.
# No source paths, hashes, GPS, model response, reviewer or arbitrary error text
# are selected. JSON label evidence is deliberately excluded at the SQL layer.
_PROGRESS_SQL = """
WITH eligible AS (
    SELECT b.batch_id,
           CASE WHEN b.processing_info->>'status' IS NULL THEN 'queued'
                WHEN b.processing_info->>'status' IN
                    ('queued','extracting','labeling','completed','validation_failed','failed','skipped')
                THEN b.processing_info->>'status' ELSE 'unknown' END AS status,
           b.review_status, b.validation_passed, b.updated_at,
           b.processing_info->>'heartbeat_at' AS heartbeat_at,
           CASE WHEN jsonb_typeof(b.processing_info->'grouping_warning') = 'array'
                THEN jsonb_array_length(b.processing_info->'grouping_warning') > 0
                ELSE (b.processing_info->'grouping_warning') = 'true'::jsonb
                END AS grouping_warning,
           COUNT(*) AS active_videos
    FROM aerial.annotation_batches b
    JOIN aerial.annotation_batch_videos bv ON bv.batch_id = b.batch_id
    JOIN aerial.videos v ON v.sha256 = bv.video_sha256 AND v.deleted_at IS NULL
    GROUP BY b.batch_id
), status_totals AS (
    SELECT status, COUNT(*) AS batches, SUM(active_videos) AS videos
    FROM eligible GROUP BY status
), ranked AS (
    SELECT e.*, ROW_NUMBER() OVER (
        PARTITION BY (status IN ('extracting','labeling'))
        ORDER BY updated_at DESC, batch_id
    ) AS position
    FROM eligible e
    WHERE status IN ('extracting','labeling','completed','validation_failed','failed','skipped')
), details AS (
    SELECT e.status, e.updated_at, jsonb_build_object(
        'batch_id', e.batch_id, 'status', e.status,
        'city', b.processing_info->>'city',
        'attempts', b.processing_info->'attempts',
        'active_videos', e.active_videos,
        'review_status', e.review_status, 'validation_passed', e.validation_passed,
        'started_at', b.processing_info->>'started_at',
        'finished_at', b.processing_info->>'finished_at',
        'heartbeat_at', e.heartbeat_at, 'updated_at', e.updated_at,
        'error_code', b.processing_info->>'error_code',
        'grouping_warning', e.grouping_warning,
        'label_count', jsonb_array_length(b.labels),
        'labels', CASE WHEN e.status = 'completed' AND e.validation_passed
                               AND e.review_status = 'pending'
                      THEN COALESCE((
                          SELECT jsonb_agg(jsonb_build_object(
                              'field', item->'field', 'state', item->'state',
                              'values', item->'values', 'scope', item->'scope'))
                          FROM jsonb_array_elements(b.labels) item
                      ), '[]'::jsonb) ELSE '[]'::jsonb END
    ) AS public_row
    FROM ranked e JOIN aerial.annotation_batches b ON b.batch_id = e.batch_id
    WHERE e.position <= CASE WHEN e.status IN ('extracting','labeling') THEN 10 ELSE 20 END
)
SELECT jsonb_build_object(
    'active_videos', (SELECT COUNT(*) FROM aerial.videos WHERE deleted_at IS NULL),
    'eligible_batches', COUNT(*), 'eligible_videos', COALESCE(SUM(active_videos), 0),
    'completed_batches', COUNT(*) FILTER (WHERE status = 'completed' AND validation_passed),
    'completed_videos', COALESCE(SUM(active_videos) FILTER (WHERE status = 'completed' AND validation_passed), 0),
    'pending_review_batches', COUNT(*) FILTER (
        WHERE status = 'completed' AND validation_passed AND review_status = 'pending'),
    'approved_batches', COUNT(*) FILTER (
        WHERE status = 'completed' AND validation_passed AND review_status = 'approved'),
    'grouping_warning_batches', COUNT(*) FILTER (WHERE grouping_warning),
    'last_heartbeat', MAX(heartbeat_at), 'latest_updated_at', MAX(updated_at),
    'working_heartbeats', COALESCE(jsonb_agg(heartbeat_at) FILTER (
        WHERE status IN ('extracting','labeling')), '[]'::jsonb),
    'status_counts', COALESCE((SELECT jsonb_agg(to_jsonb(s)) FROM status_totals s), '[]'::jsonb),
    'working', COALESCE((SELECT jsonb_agg(public_row ORDER BY updated_at DESC)
        FROM details WHERE status IN ('extracting','labeling')), '[]'::jsonb),
    'recent', COALESCE((SELECT jsonb_agg(public_row ORDER BY updated_at DESC)
        FROM details WHERE status NOT IN ('extracting','labeling')), '[]'::jsonb)
) FROM eligible
"""


@dataclass(frozen=True)
class StatusCount:
    status: str
    batches: int
    videos: int


@dataclass(frozen=True)
class CandidateLabel:
    field: str
    name: str
    state: str
    values: tuple[str, ...]
    scope: str


@dataclass(frozen=True)
class AnnotationBatch:
    batch_id: str
    city: str
    status: str
    attempts: int
    active_videos: int
    label_count: int
    review_status: str
    validation_passed: bool
    started_at: datetime | None
    finished_at: datetime | None
    heartbeat_at: datetime | None
    updated_at: datetime | None
    error_code: str
    grouping_warning: bool
    labels: tuple[CandidateLabel, ...]


@dataclass(frozen=True)
class AnnotationProgress:
    active_videos: int
    eligible_batches: int
    eligible_videos: int
    completed_batches: int
    completed_videos: int
    pending_review_batches: int
    approved_batches: int
    grouping_warning_batches: int
    stale_working_batches: int
    last_heartbeat: datetime | None
    latest_updated_at: datetime | None
    status_counts: tuple[StatusCount, ...]
    working: tuple[AnnotationBatch, ...]
    recent: tuple[AnnotationBatch, ...]

    @property
    def progress(self) -> float:
        return min(self.completed_batches / self.eligible_batches, 1.0) if self.eligible_batches else 0.0

    @property
    def attempted_batches(self) -> int:
        return sum(row.batches for row in self.status_counts
                   if row.status in ("completed", "validation_failed", "failed", "skipped"))

    @property
    def processing_progress(self) -> float:
        return min(self.attempted_batches / self.eligible_batches, 1.0) if self.eligible_batches else 0.0


def _count(value: object) -> int:
    return max(value, 0) if isinstance(value, int) and not isinstance(value, bool) else 0


def _time(value: object) -> datetime | None:
    try:
        result = value if isinstance(value, datetime) else datetime.fromisoformat(value)
        return result if result.tzinfo is not None else None
    except (ValueError, TypeError):
        return None


def _stale_working_count(raw: object) -> int:
    """Check every working heartbeat without casting untrusted text in SQL."""
    if not isinstance(raw, (list, tuple)):
        return 0
    now = datetime.now(timezone.utc)
    return sum((heartbeat := _time(value)) is None or
               (now - heartbeat).total_seconds() > 180 for value in raw)


@lru_cache(maxsize=4)
def _public_cities(path: str, modified_ns: int) -> frozenset[str]:
    collection = json.loads(Path(path).read_text(encoding="utf-8"))
    names = {row["properties"]["地名"] for row in collection["features"]}
    names.difference_update(("福州市", "抚州市"))
    names.update(("Fuzhou (Fujian)", "Fuzhou (Jiangxi)", "缺少坐标", "边界外或坐标异常"))
    return frozenset(names)


@lru_cache(maxsize=1)
def _taxonomy() -> dict[str, dict]:
    # Versioned public taxonomy is shared with the worker's strict validator.
    definitions = json.loads((_DATA / "annotation-fields-v1.json").read_text(encoding="utf-8"))
    return {row["code"]: row for row in definitions}


def _candidate_labels(raw: object) -> tuple[CandidateLabel, ...]:
    if not isinstance(raw, list):
        return ()
    definitions = _taxonomy()
    labels = []
    for item in raw[:200]:
        if not isinstance(item, dict):
            continue
        code, state, scope, values = (item.get(key) for key in ("field", "state", "scope", "values"))
        if (not isinstance(code, str) or code not in definitions
                or not isinstance(state, str) or state not in _STATES
                or not isinstance(scope, str) or scope not in _SCOPES
                or not isinstance(values, list) or len(values) > 30):
            continue
        definition = definitions[code]
        if not all(isinstance(value, str) and value in definition["values"] for value in values):
            continue
        if (len(set(values)) != len(values) or (not definition["multi"] and len(values) > 1)
                or "无" in values or (state == "observed" and not values)
                or (state != "observed" and values)
                or (definition["unobservable"] and state not in ("unknown", "unlabeled"))):
            continue
        path = definition.get("original_path")
        name = " / ".join(path[1:]) if path else definition["name"]
        labels.append(CandidateLabel(code, name, state, tuple(values), scope))
    return tuple(labels)


def _batch(raw: Mapping, cities: frozenset[str]) -> AnnotationBatch | None:
    try:
        identity = UUID(str(raw.get("batch_id"))).hex[-16:]
    except (ValueError, TypeError, AttributeError):
        return None
    status = raw.get("status")
    status = status if isinstance(status, str) and status in _STATUSES else "unknown"
    review = raw.get("review_status")
    review = review if isinstance(review, str) and review in ("pending", "approved", "rejected") else "unknown"
    valid = raw.get("validation_passed") is True
    city = raw.get("city")
    city = city if isinstance(city, str) and city in cities else "未知"
    error = raw.get("error_code")
    error = error if isinstance(error, str) and error in _ERRORS else ""
    labels = _candidate_labels(raw.get("labels")) if status == "completed" and valid and review == "pending" else ()
    return AnnotationBatch(identity, city, status, _count(raw.get("attempts")),
                           _count(raw.get("active_videos")), _count(raw.get("label_count")),
                           review, valid, *(_time(raw.get(key)) for key in
                                           ("started_at", "finished_at", "heartbeat_at", "updated_at")),
                           error, raw.get("grouping_warning") is True, labels)


def progress_snapshot(payload: Mapping) -> AnnotationProgress:
    """Validate the public projection; malformed fields never become UI text."""
    city_path = Path(os.environ.get("NDS_CITY_BOUNDARIES_PATH", str(_DATA / "prefectures-2023.geojson")))
    cities = _public_cities(str(city_path), city_path.stat().st_mtime_ns)
    status_counts = []
    raw_counts = payload.get("status_counts", ())
    for row in raw_counts if isinstance(raw_counts, (list, tuple)) else ():
        if isinstance(row, dict):
            status = row.get("status")
            status = status if isinstance(status, str) and status in _STATUSES else "unknown"
            status_counts.append(StatusCount(status, _count(row.get("batches")), _count(row.get("videos"))))
    rank = {status: index for index, status in enumerate((*_STATUSES, "unknown"))}
    status_counts.sort(key=lambda row: rank[row.status])
    def batches(key: str, limit: int) -> tuple[AnnotationBatch, ...]:
        raw = payload.get(key, ())
        if not isinstance(raw, (list, tuple)):
            return ()
        return tuple(batch for row in raw[:limit] if isinstance(row, dict)
                     if (batch := _batch(row, cities)) is not None)
    return AnnotationProgress(
        *(_count(payload.get(key)) for key in (
            "active_videos", "eligible_batches", "eligible_videos", "completed_batches",
            "completed_videos", "pending_review_batches", "approved_batches", "grouping_warning_batches")),
        _stale_working_count(payload.get("working_heartbeats")),
        _time(payload.get("last_heartbeat")), _time(payload.get("latest_updated_at")),
        tuple(status_counts), batches("working", 10), batches("recent", 20),
    )


def fetch_annotation_progress(dsn: str) -> AnnotationProgress:
    """One bounded read-only query; schema/connection errors remain unavailable."""
    with psycopg.connect(dsn, connect_timeout=5,
                         options="-c default_transaction_read_only=on -c statement_timeout=5000") as connection:
        with connection.cursor() as cursor:
            cursor.execute(_PROGRESS_SQL)
            payload = cursor.fetchone()[0]
    return progress_snapshot(payload)


@st.fragment(run_every="10s")
def render_annotation_progress(dsn: str, translate: Callable[[str], str]) -> None:
    """An independent refresh; candidate labels never enter official statistics."""
    def tr(zh: str, en: str) -> str:
        result = translate(zh)
        return en if st.session_state.get("lang") == "en" and result == zh else result

    def timestamp(value: datetime | None) -> str:
        return value.astimezone(_TIMEZONE).strftime("%Y-%m-%d %H:%M:%S") if value else "—"

    status_names = {
        "queued": tr("排队中", "Queued"), "extracting": tr("提取代表帧", "Extracting frame"),
        "labeling": tr("模型标注中", "Labeling"), "completed": tr("标注完成", "Completed"),
        "validation_failed": tr("校验失败", "Validation failed"), "failed": tr("任务失败", "Failed"),
        "skipped": tr("已跳过", "Skipped"), "unknown": tr("未知", "Unknown"),
    }
    review_names = {"pending": tr("未人工审核", "Not human-reviewed"),
                    "approved": tr("已审核", "Approved"), "rejected": tr("已退回", "Rejected"),
                    "unknown": tr("未知", "Unknown")}
    st.subheader(tr("实时标注进度", "Live annotation progress"), anchor="annotation-progress")
    st.caption(tr("每 10 秒自动刷新 · 时间为北京时间 · 统计所有关联有效视频的标注批次",
                  "Refreshes every 10 seconds · China Standard Time · All batches with active video members"))
    if not dsn:
        st.info(tr("尚未配置 PostgreSQL 数据源。", "PostgreSQL data source is not configured."))
        return
    try:
        snapshot = fetch_annotation_progress(dsn)
    except Exception:
        st.warning(tr("标注进度暂不可用，请稍后刷新。", "Annotation progress is unavailable. Please refresh later."))
        return

    st.info(tr("通过校验的自动标注结果用于标签分布统计，当前无需人工审核。标签描述代表帧观察；校验通过不等同于人工确认标签正确。",
               "Validated automatic annotations feed label distributions without a human review prerequisite. Labels describe representative frames; validation does not mean a human confirmed their accuracy."))
    metrics = (
        (tr("有效标注批次", "Eligible batches"), f"{snapshot.eligible_batches:,}"),
        (tr("校验通过批次", "Validated batches"), f"{snapshot.completed_batches:,} / {snapshot.eligible_batches:,}"),
        (tr("已生成候选批次的关联视频", "Videos linked to validated batches"), f"{snapshot.completed_videos:,} / {snapshot.active_videos:,}"),
        (tr("排队或处理中的批次", "Queued or working batches"),
         f"{sum(row.batches for row in snapshot.status_counts if row.status in ('queued', 'extracting', 'labeling')):,}"),
    )
    for column, (name, value) in zip(st.columns(4), metrics):
        column.metric(name, value)
    st.caption(tr("关联视频尚未逐一观察；动态标签仅描述代表帧中的观察结果。",
                  "Individual linked source videos have not all been observed. Dynamic labels describe only observations in the representative frame."))
    st.progress(snapshot.processing_progress,
                text=tr("本轮处理进度（含失败与跳过）", "Processing progress (includes failures and skips)")
                + f" {snapshot.attempted_batches:,} / {snapshot.eligible_batches:,} ({snapshot.processing_progress:.1%})")
    st.progress(snapshot.progress, text=tr("校验通过进度", "Validation progress") + f" {snapshot.progress:.1%}")
    st.caption(f"{tr('已关联视频', 'Linked active videos')}: {snapshot.eligible_videos:,} / {snapshot.active_videos:,}  ·  "
               f"{tr('已人工审核批次', 'Human-approved batches')}: {snapshot.approved_batches:,}  ·  "
               f"{tr('最近心跳', 'Last heartbeat')}: {timestamp(snapshot.last_heartbeat)}  ·  "
               f"{tr('库内最近更新', 'Latest database update')}: {timestamp(snapshot.latest_updated_at)}")
    if not snapshot.eligible_batches:
        st.info(tr("尚无关联有效视频的标注批次。", "There are no annotation batches with active video members."))
        return
    if snapshot.stale_working_batches:
        st.warning(tr("运行中任务缺少有效心跳或心跳超过 3 分钟未更新，请检查标注服务。",
                      "Working batches have missing or invalid heartbeats, or heartbeats older than 3 minutes. Check the annotation service.")
                   + f" ({snapshot.stale_working_batches:,})")
    if snapshot.grouping_warning_batches:
        st.warning(tr("部分采集分组需人工核查分组信息与代表帧。",
                      "Some collection groups need manual review of grouping information and representative frames.")
                   + f" ({snapshot.grouping_warning_batches:,})")
    st.dataframe(pd.DataFrame([
        {tr("处理状态", "Processing status"): status_names[row.status],
         tr("批次数", "Batches"): row.batches, tr("关联有效视频数", "Active video members"): row.videos}
        for row in snapshot.status_counts
    ]), hide_index=True, width="stretch")

    def table(rows: tuple[AnnotationBatch, ...]) -> None:
        if not rows:
            st.caption(tr("暂无记录", "No records"))
            return
        st.dataframe(pd.DataFrame([
            {tr("批次", "Batch"): row.batch_id, tr("城市", "City"): translate(row.city),
             tr("处理状态", "Processing status"): status_names[row.status],
             tr("关联有效视频数", "Active video members"): row.active_videos,
             tr("尝试次数", "Attempts"): row.attempts, tr("标签数", "Labels"): row.label_count,
             tr("人工审核", "Human review"): review_names[row.review_status],
             tr("开始时间", "Started"): timestamp(row.started_at),
             tr("完成时间", "Finished"): timestamp(row.finished_at),
             tr("心跳时间", "Heartbeat"): timestamp(row.heartbeat_at),
             tr("最近更新", "Updated"): timestamp(row.updated_at),
             tr("原因", "Reason"): tr(*_ERRORS[row.error_code]) if row.error_code else "—"}
            for row in rows
        ]), hide_index=True, width="stretch")

    st.markdown("**" + tr("正在处理的批次", "Working batches") + "**")
    table(snapshot.working)
    st.markdown("**" + tr("最近完成、失败或跳过的批次", "Recent completed, failed or skipped batches") + "**")
    table(snapshot.recent)
    candidates = [row for row in snapshot.recent if row.labels]
    if candidates:
        with st.expander(tr("自动标注结果预览", "Automatic annotation preview")):
            state_names = {"observed": tr("已观察", "Observed"), "not_observed": tr("采样帧未观察到", "Not observed in sampled frame"),
                           "unknown": tr("未知", "Unknown"), "unlabeled": tr("未标注", "Unlabeled"),
                           "not_applicable": tr("不适用", "Not applicable")}
            st.dataframe(pd.DataFrame([
                {tr("批次", "Batch"): row.batch_id, tr("城市", "City"): translate(row.city),
                 tr("标签字段", "Label field"): label.name, tr("观察状态", "Observation state"): state_names[label.state],
                 tr("候选值", "Candidate values"): "、".join(label.values) or "—",
                 tr("证据范围", "Evidence scope"): label.scope,
                 tr("人工审核", "Human review"): review_names["pending"]}
                for row in candidates for label in row.labels
            ]), hide_index=True, width="stretch")