app.py

"""Frame-extraction service — HF Space CPU worker.

GET  /video-meta       ?key=<s3_key>
  Probes video without full download: S3 HeadObject (size) + ffprobe (duration,
  fps, resolution). Returns { size_bytes, duration_s, native_fps, width, height }.
  The TS video_planner node calls this to decide target_fps before the GPU job.

POST /extract-frames   { assessment_id, video_s3_key?, frames?: [int,...], target_fps?: float }
  Seeks to each requested frame (or all if omitted), uploads raw JPEGs.
  If target_fps < native_fps, resamples the video first so frame indices match
  track.json. Fast path for interrupt-1: only the ~15 excluded frames.

POST /render           { assessment_id, video_s3_key?, track_json_s3_key,
                         bundle_s3_keys, target_fps?: float }
  Decodes all frames, draws skeleton overlay (green=qualified / red=excluded),
  stitches debug.mp4, uploads everything. For the results view.

GET  /health           → {"ok": true}

Self-contained — zero ergo_agent imports. S3 and cv2 are injected at startup.
"""
from __future__ import annotations

import json
import logging
import math
import os
import shutil
import subprocess
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Optional

import boto3
from fastapi import FastAPI, Query
from pydantic import BaseModel

logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO"),
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
)
logger = logging.getLogger("render_worker")

# ---------------------------------------------------------------------------
# Config (own — no ergo_agent.config)
# ---------------------------------------------------------------------------

BUCKET: str = os.environ["S3_BUCKET"]
_PREFIX: str = os.environ.get("S3_ASSESSMENT_PREFIX", "industrial/assessments").rstrip("/")


def _make_s3():
    return boto3.client(
        "s3",
        endpoint_url=os.environ.get("S3_ENDPOINT_URL") or None,
        region_name=os.environ.get("S3_REGION") or None,
        aws_access_key_id=os.environ["S3_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["S3_SECRET_ACCESS_KEY"],
    )


_s3 = None  # initialised lazily on first request


def _get_s3():
    global _s3
    if _s3 is None:
        _s3 = _make_s3()
    return _s3


def _assessment_prefix(assessment_id: str) -> str:
    return f"{_PREFIX}/{assessment_id}"


def _video_key(assessment_id: str, explicit: Optional[str]) -> str:
    return explicit or f"{_assessment_prefix(assessment_id)}/video/input.mp4"


# ---------------------------------------------------------------------------
# Adaptive FPS (spec: docs/VIDEO_SAMPLING_SPEC.md)
# ---------------------------------------------------------------------------

def adaptive_fps(duration_s: float, native_fps: float, size_bytes: int) -> float:
    """Return the target FPS for processing based on video characteristics.

    Never upsamples. Floor at 5fps (below this REBA temporal reasoning breaks).
    """
    # Duration-based table
    if duration_s <= 30:
        target = min(native_fps, 30.0)
    elif duration_s <= 60:
        target = 15.0
    elif duration_s <= 120:
        target = 10.0
    else:
        target = 8.0

    # Large-file penalty: >500MB AND duration >60s → reduce by 2 more
    if size_bytes > 500_000_000 and duration_s > 60:
        target = target - 2.0

    # Hard constraints
    target = max(target, 5.0)           # floor
    target = min(target, native_fps)    # never upsample
    return float(target)


# ---------------------------------------------------------------------------
# Video resampling (ensures frame indices match track.json)
# ---------------------------------------------------------------------------

def _resample_video(src: str, target_fps: float, out_dir: str) -> str:
    """Re-encode video at target_fps. Returns path to resampled file."""
    out_path = os.path.join(out_dir, "resampled.mp4")
    cmd = [
        "ffmpeg", "-y", "-loglevel", "error",
        "-i", src,
        "-vf", f"fps={target_fps:.6g}",
        "-an",
        "-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
        "-pix_fmt", "yuv420p",
        out_path,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(
            f"ffmpeg resample failed (exit {proc.returncode}): "
            f"{proc.stderr.strip() or '<no stderr>'}"
        )
    return out_path


def _transcode_h264(src: str, dst: str) -> None:
    """Re-encode src (any codec) to H.264 in a browser-compatible MP4."""
    cmd = [
        "ffmpeg", "-y", "-loglevel", "error",
        "-i", src,
        "-c:v", "libx264", "-preset", "veryfast", "-crf", "23",
        "-pix_fmt", "yuv420p",  # required for Safari + broad browser compat
        "-movflags", "+faststart",  # moov atom at front — playback starts immediately
        "-an",
        dst,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(
            f"ffmpeg H.264 transcode failed (exit {proc.returncode}): "
            f"{proc.stderr.strip() or '<no stderr>'}"
        )


def _probe_video(local_path: str) -> dict:
    """Run ffprobe on a local file; returns {duration_s, native_fps, width, height}."""
    cmd = [
        "ffprobe", "-v", "quiet", "-print_format", "json",
        "-show_streams", "-show_format", local_path,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"ffprobe failed: {proc.stderr.strip()}")
    info = json.loads(proc.stdout)
    stream = next((s for s in info.get("streams", []) if "width" in s), {})
    duration_s = float(info.get("format", {}).get("duration", 0))
    r_fps = stream.get("r_frame_rate", "0/1")
    num, den = (int(x) for x in r_fps.split("/"))
    native_fps = num / den if den else 0.0
    return {
        "duration_s": duration_s,
        "native_fps": native_fps,
        "width": stream.get("width", 0),
        "height": stream.get("height", 0),
    }


# ---------------------------------------------------------------------------
# Skeleton — per-edge severity coloring
# ---------------------------------------------------------------------------

# Each edge: (joint_a, joint_b, part_name)
# part_name maps to rec["parts"][part]["score"] in L1.
# Joints: NECK=69, NOSE=0, L_SH=5, R_SH=6, L_EL=7, R_EL=8,
#         L_HIP=9, R_HIP=10, L_KNEE=11, R_KNEE=12, L_ANK=13, R_ANK=14,
#         L_WRIST=62, R_WRIST=41
_EDGES = [
    (69,  5, "trunk"),      # neck → left shoulder
    (69,  6, "trunk"),      # neck → right shoulder
    ( 5,  7, "upper_arm"),  # left shoulder → left elbow
    ( 7, 62, "lower_arm"),  # left elbow → left wrist
    ( 6,  8, "upper_arm"),  # right shoulder → right elbow
    ( 8, 41, "lower_arm"),  # right elbow → right wrist
    ( 5,  9, "trunk"),      # left shoulder → left hip
    ( 6, 10, "trunk"),      # right shoulder → right hip
    ( 9, 10, "trunk"),      # left hip → right hip
    ( 9, 11, "legs"),       # left hip → left knee
    (11, 13, "legs"),       # left knee → left ankle
    (10, 12, "legs"),       # right hip → right knee
    (12, 14, "legs"),       # right knee → right ankle
    ( 0, 69, "neck"),       # nose → neck
]

# BGR — score thresholds match REBA risk bands: negligible=1, low=2-3, medium=4-7, high=8-10, very_high=11+
_SCORE_COLOR = [
    (1,  (0, 200,   0)),   # 1        negligible — green
    (3,  (0, 200, 160)),   # 2–3      low        — teal
    (7,  (0, 180, 255)),   # 4–7      medium     — amber
    (10, (0,  80, 255)),   # 8–10     high       — orange
    (99, (40,  40, 220)),  # 11+      very_high  — red
]

_EXCLUDED_COLOR = (80, 80, 80)  # grey — disqualified frame, don't read score


def _part_color(score: int) -> tuple:
    for threshold, color in _SCORE_COLOR:
        if score <= threshold:
            return color
    return _SCORE_COLOR[-1][1]


def overlay_segments(
    q: list, qv: list, w: int, h: int, parts: dict, vis_floor: float = 0.3
) -> list:
    """Pure: normalised 2D joints → colored pixel line segments.

    Each edge is colored by the REBA score of its body part. Returns
    [((x1,y1),(x2,y2), color), ...] in pixel coords.
    """
    segs = []
    for a, b, part in _EDGES:
        if a >= len(q) or b >= len(q):
            continue
        if qv[a] < vis_floor or qv[b] < vis_floor:
            continue
        pa = (int(q[a][0] * w), int(q[a][1] * h))
        pb = (int(q[b][0] * w), int(q[b][1] * h))
        score = (parts.get(part) or {}).get("score") or 1
        segs.append((pa, pb, _part_color(score)))
    return segs


# ---------------------------------------------------------------------------
# S3 helpers
# ---------------------------------------------------------------------------

def _put_frame(rbase: str, filename: str, jpg_bytes: bytes) -> None:
    _get_s3().put_object(
        Bucket=BUCKET,
        Key=f"{rbase}/frames/{filename}",
        Body=jpg_bytes,
        ContentType="image/jpeg",
    )


def _put_video(rbase: str, local_path: str) -> None:
    if not os.path.exists(local_path):
        return
    _get_s3().put_object(
        Bucket=BUCKET,
        Key=f"{rbase}/debug.mp4",
        Body=open(local_path, "rb").read(),
        ContentType="video/mp4",
    )


def _open_video(out_dir: str, s3_key: str, target_fps: Optional[float]):
    """Download video from S3, optionally resample, return (cv2_cap, vid_path)."""
    import cv2  # lazy
    vid_path = os.path.join(out_dir, "v.mp4")
    _get_s3().download_file(BUCKET, s3_key, vid_path)

    if target_fps is not None:
        # probe native fps to decide whether resampling is actually needed
        probe = _probe_video(vid_path)
        if target_fps < probe["native_fps"] - 0.01:
            vid_path = _resample_video(vid_path, target_fps, out_dir)

    return cv2.VideoCapture(vid_path), vid_path


# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------

class ExtractReq(BaseModel):
    assessment_id: str
    video_s3_key: Optional[str] = None
    frames: Optional[List[int]] = None  # None → all frames
    target_fps: Optional[float] = None  # None → use native fps as-is


class RenderReq(BaseModel):
    assessment_id: str
    video_s3_key: Optional[str] = None
    track_json_s3_key: str
    bundle_s3_keys: dict
    target_fps: Optional[float] = None


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI()


@app.get("/health")
def health() -> dict:
    return {"ok": True}


@app.get("/video-meta")
def video_meta(key: str = Query(..., description="S3 key of the video")) -> dict:
    """Probe video metadata without full download.

    S3 HeadObject → size_bytes. Downloads a small header via ffprobe for
    duration/fps/resolution. The TS video_planner node calls this once before
    the GPU job to decide target_fps.
    """
    logger.info("video-meta: probing key=%s", key)
    t0 = time.monotonic()
    head = _get_s3().head_object(Bucket=BUCKET, Key=key)
    size_bytes = head["ContentLength"]

    # Download just enough for ffprobe: stream from S3 to a temp file
    # (ffprobe only needs the container header, but boto3 download_file
    #  fetches the whole thing — acceptable for metadata since it's async/cheap)
    out = tempfile.mkdtemp(prefix="meta_")
    try:
        vid_path = os.path.join(out, "v.mp4")
        _get_s3().download_file(BUCKET, key, vid_path)
        probe = _probe_video(vid_path)
    finally:
        shutil.rmtree(out, ignore_errors=True)

    result = {
        "size_bytes": size_bytes,
        "duration_s": probe["duration_s"],
        "native_fps": probe["native_fps"],
        "width": probe["width"],
        "height": probe["height"],
    }
    logger.info(
        "video-meta: done in %.1fs — %.1fs @ %.2ffps %dx%d (%.1fMB)",
        time.monotonic() - t0,
        probe["duration_s"], probe["native_fps"],
        probe["width"], probe["height"],
        size_bytes / 1e6,
    )
    return result


def _extract_specific(cv2, cap, pool, rbase: str, wanted: list) -> int:
    """Seek to each requested frame index and upload. Returns count uploaded."""
    futures = []
    n = 0
    for idx in wanted:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ok, frame = cap.read()
        if not ok:
            logger.warning("extract-frames: seek to frame %d failed", idx)
            continue
        ok2, jpg = cv2.imencode(".jpg", frame)
        if ok2:
            futures.append(pool.submit(_put_frame, rbase, f"f{idx:04d}.jpg", jpg.tobytes()))
            n += 1
    for fut in as_completed(futures):
        fut.result()
    return n


def _extract_all(cv2, cap, pool, rbase: str) -> int:
    """Read every frame sequentially and upload. Returns count uploaded."""
    futures = []
    n = 0
    i = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        ok2, jpg = cv2.imencode(".jpg", frame)
        if ok2:
            futures.append(pool.submit(_put_frame, rbase, f"f{i:04d}.jpg", jpg.tobytes()))
            n += 1
        i += 1
    for fut in as_completed(futures):
        fut.result()
    return n


@app.post("/extract-frames")
def do_extract(req: ExtractReq) -> dict:
    rbase = f"{_assessment_prefix(req.assessment_id)}/render"
    n_requested = len(req.frames) if req.frames is not None else None
    logger.info(
        "extract-frames: assessment=%s frames=%s target_fps=%s",
        req.assessment_id,
        f"{n_requested} specific" if n_requested is not None else "all",
        req.target_fps,
    )
    t0 = time.monotonic()
    out = tempfile.mkdtemp(prefix="extract_")
    try:
        t_dl = time.monotonic()
        cap, _ = _open_video(out, _video_key(req.assessment_id, req.video_s3_key), req.target_fps)
        logger.info("extract-frames: video ready in %.1fs", time.monotonic() - t_dl)
        import cv2  # lazy — stubbed in tests
        wanted = sorted(set(req.frames)) if req.frames is not None else None
        try:
            with ThreadPoolExecutor(max_workers=8) as pool:
                n = (
                    _extract_specific(cv2, cap, pool, rbase, wanted)
                    if wanted is not None
                    else _extract_all(cv2, cap, pool, rbase)
                )
        finally:
            cap.release()
    finally:
        shutil.rmtree(out, ignore_errors=True)

    logger.info(
        "extract-frames: done — %d frames uploaded to %s in %.1fs",
        n, rbase, time.monotonic() - t0,
    )
    return {"render_s3_prefix": rbase, "n_frames": n}


def _overlay_frame(cv2, frame: object, i: int, track_frames: list, l1: list, w: int, h: int) -> int:
    """Draw per-part severity skeleton + frame score label. Returns segments drawn."""
    tf = track_frames[i] if i < len(track_frames) else {}
    rec = l1[i] if i < len(l1) else {}
    q = tf.get("q") or []
    qv = tf.get("qv") or []
    if not (q and qv):
        return 0

    qualified = rec.get("qualification", {}).get("qualified", True)
    parts = rec.get("parts") or {}
    final = rec.get("final")

    if not qualified:
        # Excluded frame: draw everything grey so it's visible but clearly flagged
        segs = overlay_segments(q, qv, w, h, {})
        for pa, pb, _ in segs:
            cv2.line(frame, pa, pb, _EXCLUDED_COLOR, 2)
        label = f"f{i:04d} [excl]"
        cv2.putText(frame, label, (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, _EXCLUDED_COLOR, 2)
        return len(segs)

    segs = overlay_segments(q, qv, w, h, parts)
    for pa, pb, color in segs:
        cv2.line(frame, pa, pb, color, 2)

    # Frame label: index + overall REBA score colored by frame risk
    frame_score = final or 1
    label_color = _part_color(frame_score)
    risk = rec.get("risk", "")
    label = f"f{i:04d}  REBA {frame_score} {risk}"
    cv2.putText(frame, label, (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, label_color, 2)
    return len(segs)


@app.post("/render")
def do_render(req: RenderReq) -> dict:
    rbase = f"{_assessment_prefix(req.assessment_id)}/render"
    logger.info(
        "render: assessment=%s track=%s target_fps=%s",
        req.assessment_id, req.track_json_s3_key, req.target_fps,
    )
    t0 = time.monotonic()
    out = tempfile.mkdtemp(prefix="render_")
    try:
        t_dl = time.monotonic()
        cap, _ = _open_video(out, _video_key(req.assessment_id, req.video_s3_key), req.target_fps)
        import cv2  # lazy — stubbed in tests

        track_raw = _get_s3().get_object(Bucket=BUCKET, Key=req.track_json_s3_key)["Body"].read()
        track = json.loads(track_raw)
        track_frames = track.get("frames") or []

        l1_raw = _get_s3().get_object(Bucket=BUCKET, Key=req.bundle_s3_keys["l1"])["Body"].read()
        l1 = [json.loads(line) for line in l1_raw.decode().splitlines() if line.strip()]
        logger.info(
            "render: data loaded in %.1fs — %d track frames, %d l1 records",
            time.monotonic() - t_dl, len(track_frames), len(l1),
        )

        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS) or 15.0
        logger.info("render: video %dx%d @ %.2ffps", w, h, fps)

        # Write frames into a raw mp4v intermediate — OpenCV's avc1 is unavailable
        # on most Linux builds. We transcode to H.264 with ffmpeg afterward so
        # browsers can play the result in a <video> tag.
        raw_path = os.path.join(out, "raw.mp4")
        debug_path = os.path.join(out, "debug.mp4")
        vw = cv2.VideoWriter(
            raw_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
        )

        futures = []
        i = 0
        n_overlaid = 0
        t_encode = time.monotonic()
        try:
            with ThreadPoolExecutor(max_workers=8) as pool:
                while True:
                    ok, frame = cap.read()
                    if not ok:
                        break
                    n_overlaid += _overlay_frame(cv2, frame, i, track_frames, l1, w, h) > 0
                    ok2, jpg = cv2.imencode(".jpg", frame)
                    if ok2:
                        futures.append(
                            pool.submit(_put_frame, rbase, f"f{i:04d}.jpg", jpg.tobytes())
                        )
                    vw.write(frame)
                    i += 1

                for fut in as_completed(futures):
                    fut.result()
        finally:
            cap.release()
            vw.release()

        logger.info(
            "render: encoded %d frames (%d overlaid) in %.1fs; transcoding to H.264",
            i, n_overlaid, time.monotonic() - t_encode,
        )
        t_xcode = time.monotonic()
        _transcode_h264(raw_path, debug_path)
        logger.info("render: H.264 transcode done in %.1fs", time.monotonic() - t_xcode)

        t_up = time.monotonic()
        _put_video(rbase, debug_path)
        logger.info("render: debug.mp4 uploaded in %.1fs → %s", time.monotonic() - t_up, rbase)

    finally:
        shutil.rmtree(out, ignore_errors=True)

    logger.info("render: total %.1fs for assessment=%s", time.monotonic() - t0, req.assessment_id)
    return {"render_s3_prefix": rbase}