"""Frame image storage — save/load to S3/MinIO as JPEGs."""

from __future__ import annotations

import logging
import os
import tempfile

import numpy as np
from PIL import Image

from detect.models import Frame

logger = logging.getLogger(__name__)

BUCKET = os.environ.get("S3_BUCKET", "mpr")
CHECKPOINT_PREFIX = "checkpoints"


def save_frames(job_id: str, frames: list[Frame]) -> dict[int, str]:
    """
    Save frame images to S3 as JPEGs.

    Returns manifest: {sequence: s3_key}
    """
    from core.storage.s3 import upload_file

    manifest = {}

    for frame in frames:
        key = f"{CHECKPOINT_PREFIX}/{job_id}/frames/{frame.sequence}.jpg"

        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            img = Image.fromarray(frame.image)
            img.save(tmp, format="JPEG", quality=85)
            tmp_path = tmp.name

        try:
            upload_file(tmp_path, BUCKET, key)
        finally:
            os.unlink(tmp_path)

        manifest[frame.sequence] = key

    logger.info("Saved %d frames to s3://%s/%s/%s/frames/",
                len(frames), BUCKET, CHECKPOINT_PREFIX, job_id)
    return manifest


def load_frames(manifest: dict[int, str], frame_metadata: list[dict]) -> list[Frame]:
    """
    Load frame images from S3 and reconstitute Frame objects.

    frame_metadata: list of dicts with sequence, chunk_id, timestamp, perceptual_hash.
    """
    from core.storage.s3 import download_to_temp

    meta_map = {m["sequence"]: m for m in frame_metadata}
    frames = []

    for seq, key in manifest.items():
        tmp_path = download_to_temp(BUCKET, key)
        try:
            img = Image.open(tmp_path).convert("RGB")
            image_array = np.array(img)
        finally:
            os.unlink(tmp_path)

        meta = meta_map.get(seq, {})
        frame = Frame(
            sequence=seq,
            chunk_id=meta.get("chunk_id", 0),
            timestamp=meta.get("timestamp", 0.0),
            image=image_array,
            perceptual_hash=meta.get("perceptual_hash", ""),
        )
        frames.append(frame)

    frames.sort(key=lambda f: f.sequence)
    return frames