mediaproc/detect/checkpoint/storage.py

"""
Checkpoint storage — Timeline + Checkpoint (tree of snapshots).

Timeline: frame sequence from source video (frames in MinIO)
Checkpoint: snapshot of pipeline state (stage outputs as JSONB in Postgres)
            parent_id forms a tree — multiple children = different config tries
"""

from __future__ import annotations

import logging
from uuid import UUID

from .frames import save_frames, load_frames, CHECKPOINT_PREFIX

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Timeline
# ---------------------------------------------------------------------------

def create_timeline(
    source_video: str,
    profile_name: str,
    frames: list,
    fps: float = 2.0,
    source_asset_id: UUID | None = None,
) -> tuple[str, str]:
    """
    Create a timeline from frames. Uploads frame images to MinIO,
    creates Timeline + root Checkpoint in Postgres.

    Returns (timeline_id, checkpoint_id).
    """
    from core.db.detect import create_timeline as db_create_timeline
    from core.db.detect import save_checkpoint

    # Create timeline
    timeline = db_create_timeline(
        source_video=source_video,
        profile_name=profile_name,
        source_asset_id=source_asset_id,
        fps=fps,
    )
    tid = str(timeline.id)

    # Upload frames to MinIO
    manifest = save_frames(tid, frames)

    # Store frame metadata on the timeline
    frames_meta = [
        {
            "sequence": f.sequence,
            "chunk_id": getattr(f, "chunk_id", 0),
            "timestamp": f.timestamp,
            "perceptual_hash": getattr(f, "perceptual_hash", ""),
        }
        for f in frames
    ]

    timeline.frames_prefix = f"{CHECKPOINT_PREFIX}/{tid}/frames/"
    timeline.frames_manifest = {str(k): v for k, v in manifest.items()}
    timeline.frames_meta = frames_meta

    from core.db.connection import get_session
    with get_session() as session:
        session.add(timeline)
        session.commit()

    # Create root checkpoint (no parent, no stage outputs yet)
    checkpoint = save_checkpoint(
        timeline_id=timeline.id,
        parent_id=None,
        stage_outputs={},
        stats={"frames_extracted": len(frames)},
    )

    logger.info("Timeline created: %s (%d frames, root checkpoint %s)",
                tid, len(frames), checkpoint.id)
    return tid, str(checkpoint.id)


def get_timeline_frames(timeline_id: str) -> list:
    """Load frames from a timeline (from MinIO) as Frame objects."""
    from core.db.detect import get_timeline

    timeline = get_timeline(timeline_id)
    if not timeline:
        raise ValueError(f"Timeline not found: {timeline_id}")

    raw_manifest = timeline.frames_manifest or {}
    manifest = {int(k): v for k, v in raw_manifest.items()}
    frame_metadata = timeline.frames_meta or []

    return load_frames(manifest, frame_metadata)


def get_timeline_frames_b64(timeline_id: str) -> list[dict]:
    """Load frames as base64 JPEG (lightweight, no numpy)."""
    from core.db.detect import get_timeline
    from .frames import load_frames_b64

    timeline = get_timeline(timeline_id)
    if not timeline:
        raise ValueError(f"Timeline not found: {timeline_id}")

    raw_manifest = timeline.frames_manifest or {}
    manifest = {int(k): v for k, v in raw_manifest.items()}
    frame_metadata = timeline.frames_meta or []

    return load_frames_b64(manifest, frame_metadata)


# ---------------------------------------------------------------------------
# Checkpoint
# ---------------------------------------------------------------------------

def save_stage_output(
    timeline_id: str,
    parent_checkpoint_id: str | None,
    stage_name: str,
    output_json: dict,
    config_overrides: dict | None = None,
    stats: dict | None = None,
    is_scenario: bool = False,
    scenario_label: str = "",
) -> str:
    """
    Save a stage's output as a new checkpoint (child of parent).

    Carries forward stage outputs from parent + adds the new one.
    Returns the new checkpoint ID.
    """
    from core.db.detect import get_checkpoint, save_checkpoint

    # Carry forward from parent
    parent_outputs = {}
    parent_stats = {}
    parent_config = {}
    if parent_checkpoint_id:
        parent = get_checkpoint(parent_checkpoint_id)
        if parent:
            parent_outputs = dict(parent.stage_outputs or {})
            parent_stats = dict(parent.stats or {})
            parent_config = dict(parent.config_overrides or {})

    # Add new stage output
    stage_outputs = {**parent_outputs, stage_name: output_json}

    # Merge stats and config
    merged_stats = {**parent_stats, **(stats or {})}
    merged_config = {**parent_config, **(config_overrides or {})}

    checkpoint = save_checkpoint(
        timeline_id=timeline_id,
        parent_id=parent_checkpoint_id,
        stage_outputs=stage_outputs,
        config_overrides=merged_config,
        stats=merged_stats,
        is_scenario=is_scenario,
        scenario_label=scenario_label,
    )

    logger.info("Checkpoint saved: %s (timeline %s, stage %s, parent %s)",
                checkpoint.id, timeline_id, stage_name, parent_checkpoint_id)
    return str(checkpoint.id)


def load_stage_output(checkpoint_id: str, stage_name: str) -> dict | None:
    """Load a stage's output from a checkpoint."""
    from core.db.detect import get_checkpoint

    checkpoint = get_checkpoint(checkpoint_id)
    if not checkpoint:
        return None

    return (checkpoint.stage_outputs or {}).get(stage_name)