"""Visualization overlay endpoints.

Generates skeleton overlay videos from H5 pose data.
Uses the same job pattern as extraction.py (in-memory dict + thread lock + BackgroundTasks).
"""

from __future__ import annotations

import logging
import threading
import uuid
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, HTTPException

from sltk.api.dependencies import validate_path
from sltk.api.models import OverlayRequest, VizJob

router = APIRouter(prefix="/api/visualization", tags=["visualization"])
logger = logging.getLogger(__name__)

# ============================================================================
# Job Storage
# ============================================================================

_jobs: dict[str, VizJob] = {}
_jobs_lock = threading.Lock()


def _get_job(job_id: str) -> VizJob | None:
    with _jobs_lock:
        return _jobs.get(job_id)


def _set_job(job_id: str, job: VizJob) -> None:
    with _jobs_lock:
        _jobs[job_id] = job


def _get_video_frame_count(path: Path) -> int:
    """Get frame count from a video file. Returns 0 on failure."""
    try:
        import cv2

        cap = cv2.VideoCapture(str(path))
        count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()
        return max(0, count)
    except Exception:
        return 0


def _output_path_for(h5_path: Path, viz_type: str) -> Path:
    """Compute the overlay output path convention."""
    overlays_dir = h5_path.parent / "overlays"
    return overlays_dir / f"{h5_path.stem}_{viz_type}_skeleton.mp4"


# ============================================================================
# Overlay Generation Worker
# ============================================================================


def _run_overlay_job(
    job_id: str, video_path: Path, h5_path: Path, output_path: Path, viz_type: str
) -> None:
    """Background worker that generates the overlay video."""
    from sltk.visualization.pipeline import generate_overlay_video

    def progress_cb(current: int, total: int) -> None:
        with _jobs_lock:
            if job_id in _jobs:
                _jobs[job_id].current_frame = current
                _jobs[job_id].total_frames = total

    try:
        with _jobs_lock:
            if job_id in _jobs:
                _jobs[job_id].status = "running"

        generate_overlay_video(
            video_path=video_path,
            h5_path=h5_path,
            output_path=output_path,
            viz_type=viz_type,
            progress_callback=progress_cb,
        )

        with _jobs_lock:
            if job_id in _jobs:
                _jobs[job_id].status = "completed"
                _jobs[job_id].output_path = str(output_path)

    except Exception as exc:
        logger.exception("Overlay generation failed for job %s", job_id)
        with _jobs_lock:
            if job_id in _jobs:
                _jobs[job_id].status = "failed"
                _jobs[job_id].error = str(exc)


# ============================================================================
# Endpoints
# ============================================================================


@router.post("/generate", response_model=VizJob)
async def generate_overlay(request: OverlayRequest, background_tasks: BackgroundTasks) -> VizJob:
    """Start overlay video generation.

    Returns cached result immediately if overlay exists and H5 hasn't changed.
    """
    video_path = validate_path(request.video_path, must_exist=True)
    h5_path = validate_path(request.h5_path, must_exist=True)
    output_path = _output_path_for(h5_path, request.viz_type)

    # Check cache: overlay exists and is newer than H5
    if not request.force and output_path.exists():
        overlay_mtime = output_path.stat().st_mtime
        h5_mtime = h5_path.stat().st_mtime
        if overlay_mtime >= h5_mtime:
            total = _get_video_frame_count(output_path)
            return VizJob(
                job_id="cached",
                status="completed",
                output_path=str(output_path),
                total_frames=total,
                current_frame=total,
            )

    # Atomic dedup check + job creation under same lock
    with _jobs_lock:
        for jid, job in _jobs.items():
            if job.output_path == str(output_path) and job.status in ("pending", "running"):
                return job

        # Create new job inside the lock to prevent TOCTOU race
        job_id = str(uuid.uuid4())[:8]
        job = VizJob(job_id=job_id, status="pending", output_path=str(output_path))
        _jobs[job_id] = job

    background_tasks.add_task(
        _run_overlay_job, job_id, video_path, h5_path, output_path, request.viz_type
    )

    return job


@router.get("/status/{job_id}", response_model=VizJob)
async def get_status(job_id: str) -> VizJob:
    """Poll overlay generation progress."""
    job = _get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
    return job


@router.get("/check")
async def check_overlay(video_path: str, h5_path: str, viz_type: str) -> dict:
    """Check if a cached overlay exists for a video + viz_type.

    Returns:
        {exists: bool, output_path: str | None}
    """
    try:
        h5 = validate_path(h5_path, must_exist=True)
    except HTTPException:
        return {"exists": False, "output_path": None}

    output = _output_path_for(h5, viz_type)
    if output.exists():
        h5_mtime = h5.stat().st_mtime
        overlay_mtime = output.stat().st_mtime
        if overlay_mtime >= h5_mtime:
            return {"exists": True, "output_path": str(output)}
    return {"exists": False, "output_path": None}
