"""
Extraction endpoints for SLTK API.

Handles pose extraction pipeline jobs.
"""

from __future__ import annotations

import logging
import os
import threading
import uuid
from datetime import datetime, timedelta
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, HTTPException, Request

from sltk.api.dependencies import rate_limit, validate_path
from sltk.api.models import (
    ExtractionConfig,
    ExtractionJob,
    ExtractionRequest,
    StageStatus,
)

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

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

# In-memory job storage (use Redis/DB in production)
# Thread safety: Use _jobs_lock when accessing jobs or job_logs
jobs: dict[str, ExtractionJob] = {}
_jobs_lock = threading.Lock()

# In-memory log storage for extraction jobs
job_logs: dict[str, list[str]] = {}
_logs_lock = threading.Lock()

# Job cleanup configuration
JOB_TTL_HOURS = 24  # How long to keep completed jobs
MAX_JOBS = 1000  # Maximum number of jobs to keep
JOB_CLEANUP_INTERVAL_SECONDS = 3600  # Run cleanup every hour


def add_job_log(job_id: str, message: str, level: str = "info"):
    """Add a log message for a job. Thread-safe."""
    timestamp = datetime.now().strftime("%H:%M:%S")
    log_entry = f"[{timestamp}] [{level.upper()}] {message}"
    with _logs_lock:
        if job_id not in job_logs:
            job_logs[job_id] = []
        job_logs[job_id].append(log_entry)
        # Keep only last 1000 log entries
        if len(job_logs[job_id]) > 1000:
            job_logs[job_id] = job_logs[job_id][-1000:]


def cleanup_old_jobs() -> int:
    """
    Remove jobs older than TTL and enforce max job limit.

    This function is thread-safe and uses the _jobs_lock and _logs_lock
    to prevent race conditions when cleaning up jobs and logs.

    Returns:
        Number of jobs removed.
    """
    now = datetime.now()
    ttl_cutoff = now - timedelta(hours=JOB_TTL_HOURS)

    jobs_to_remove: list[str] = []

    # Find jobs to remove based on TTL
    with _jobs_lock:
        for job_id, job in jobs.items():
            # Remove completed/failed/cancelled jobs older than TTL
            if job.status in ("completed", "failed", "cancelled"):
                if job.completed_at:
                    try:
                        completed_time = datetime.fromisoformat(job.completed_at)
                        if completed_time < ttl_cutoff:
                            jobs_to_remove.append(job_id)
                    except ValueError:
                        # Invalid timestamp, mark for removal
                        jobs_to_remove.append(job_id)

        # Remove old jobs
        for job_id in jobs_to_remove:
            jobs.pop(job_id, None)

        # If still over limit, remove oldest completed jobs
        if len(jobs) > MAX_JOBS:
            completed_jobs = [
                (jid, j)
                for jid, j in jobs.items()
                if j.status in ("completed", "failed", "cancelled")
            ]
            # Sort by completed_at (oldest first), handling None
            completed_jobs.sort(key=lambda x: x[1].completed_at or "")

            excess = len(jobs) - MAX_JOBS
            for job_id, _ in completed_jobs[:excess]:
                jobs.pop(job_id, None)
                if job_id not in jobs_to_remove:
                    jobs_to_remove.append(job_id)

    # Clean up logs for removed jobs
    with _logs_lock:
        for job_id in jobs_to_remove:
            job_logs.pop(job_id, None)

    return len(jobs_to_remove)


def _get_job(job_id: str) -> ExtractionJob | None:
    """Get a job by ID with thread safety."""
    with _jobs_lock:
        return jobs.get(job_id)


def _update_job(job_id: str, **kwargs) -> None:
    """Update job fields with thread safety."""
    with _jobs_lock:
        if job_id in jobs:
            job = jobs[job_id]
            for key, value in kwargs.items():
                setattr(job, key, value)


class ExtractionCancelled(Exception):
    """Exception raised when extraction is cancelled."""

    pass


def check_job_cancelled(job_id: str) -> bool:
    """Check if a job has been cancelled. Returns True if cancelled. Thread-safe."""
    with _jobs_lock:
        if job_id in jobs and jobs[job_id].status == "cancelled":
            return True
        return False


def raise_if_cancelled(job_id: str):
    """Raise ExtractionCancelled if the job has been cancelled."""
    if check_job_cancelled(job_id):
        raise ExtractionCancelled(f"Job {job_id} was cancelled by user")


def cleanup_gpu_memory():
    """
    Force cleanup of GPU memory after job cancellation or completion.

    This helps free VRAM that may be held by extractors when jobs are
    cancelled mid-extraction.
    """
    import gc

    gc.collect()

    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
    except ImportError:
        pass  # torch not installed


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


@router.post("/start", response_model=ExtractionJob)
@rate_limit("5/minute")
async def start_extraction(
    request: Request, body: ExtractionRequest, background_tasks: BackgroundTasks
):
    """
    Start an extraction pipeline job.

    This runs in the background and can be monitored via /api/extraction/status/{job_id}.

    Rate limited to 5 requests per minute (extraction is very expensive).
    """
    # Security: Validate all paths are within allowed directories
    input_path = validate_path(body.input_path, must_exist=True)
    output_root = validate_path(body.output_root, must_exist=False, allow_creation=True)

    # Also validate any specific video paths if provided
    if body.video_paths:
        for video_path in body.video_paths:
            validate_path(video_path, must_exist=True)

    # Create output directory (already validated)
    output_root.mkdir(parents=True, exist_ok=True)

    # Create job
    job_id = str(uuid.uuid4())[:8]
    job = ExtractionJob(
        job_id=job_id,
        status="pending",
        created_at=datetime.now().isoformat(),
        input_path=str(input_path),
        output_root=str(output_root),
        config=body.config,
    )
    with _jobs_lock:
        jobs[job_id] = job

    # Start background extraction
    background_tasks.add_task(
        run_extraction_pipeline,
        job_id,
        input_path,
        output_root,
        body.config,
        body.recursive,
        body.video_paths,
    )

    return job


@router.get("/status/{job_id}", response_model=ExtractionJob)
async def get_extraction_status(job_id: str):
    """Get status of an extraction job."""
    with _jobs_lock:
        if job_id not in jobs:
            raise HTTPException(404, f"Job not found: {job_id}")
        return jobs[job_id]


@router.post("/cancel/{job_id}")
async def cancel_extraction(job_id: str):
    """Cancel an extraction job and free GPU resources."""
    with _jobs_lock:
        if job_id not in jobs:
            raise HTTPException(404, f"Job not found: {job_id}")

        job = jobs[job_id]
        if job.status == "running":
            job.status = "cancelled"
    add_job_log(job_id, "Cancellation requested - stopping at next checkpoint", "warning")

    # Schedule GPU cleanup to run after the extraction threads have had a chance to stop
    # This runs in a separate thread to not block the response
    def delayed_cleanup():
        import time

        time.sleep(2)  # Wait for extraction threads to notice cancellation and exit
        cleanup_gpu_memory()
        add_job_log(job_id, "GPU memory cleanup completed", "info")

    cleanup_thread = threading.Thread(target=delayed_cleanup, daemon=True)
    cleanup_thread.start()

    return {"status": "cancelled", "job_id": job_id}


@router.get("/jobs", response_model=list[ExtractionJob])
async def list_extraction_jobs():
    """List all extraction jobs."""
    with _jobs_lock:
        return list(jobs.values())


@router.get("/logs/{job_id}")
async def get_extraction_logs(job_id: str, offset: int = 0, limit: int = 100):
    """Get logs for an extraction job."""
    with _jobs_lock:
        if job_id not in jobs:
            raise HTTPException(404, f"Job not found: {job_id}")

    with _logs_lock:
        logs = job_logs.get(job_id, [])
        total = len(logs)
        result_logs = logs[offset : offset + limit]

    return {
        "job_id": job_id,
        "logs": result_logs,
        "total": total,
        "offset": offset,
        "has_more": offset + limit < total,
    }


# ============================================================================
# Pipeline Implementation
# ============================================================================


async def run_extraction_pipeline(
    job_id: str,
    input_path: Path,
    output_root: Path,
    config: ExtractionConfig,
    recursive: bool,
    video_paths: list[str] | None = None,
):
    """
    Run the full extraction pipeline.

    This is executed as a background task.
    """
    job = _get_job(job_id)
    if job is None:
        return
    _update_job(job_id, status="running", started_at=datetime.now().isoformat())

    add_job_log(job_id, "Starting extraction pipeline", "info")
    add_job_log(job_id, f"Input path: {input_path}", "info")
    add_job_log(job_id, f"Output root: {output_root}", "info")
    config_msg = (
        f"Config: mediapipe={config.enable_mediapipe}, "
        f"wilor={config.enable_wilor}, nlf={config.enable_nlf}"
    )
    add_job_log(job_id, config_msg, "info")

    try:
        # Use specific video paths if provided, otherwise discover
        if video_paths:
            video_files = [Path(p) for p in video_paths]
            add_job_log(job_id, f"Using {len(video_files)} specified video paths", "info")
        elif input_path.is_file():
            video_files = [input_path]
            add_job_log(job_id, "Processing single video file", "info")
        else:
            extensions = {".mp4", ".mov", ".avi", ".mkv"}
            video_files = []

            add_job_log(job_id, "Scanning directory for videos...", "info")

            def handle_error(err: OSError):
                add_job_log(job_id, f"Permission denied: {err.filename}", "warning")

            if recursive:
                for dirpath, dirnames, filenames in os.walk(input_path, onerror=handle_error):
                    dir_path = Path(dirpath)
                    for filename in filenames:
                        if any(filename.lower().endswith(ext) for ext in extensions):
                            video_files.append(dir_path / filename)
            else:
                try:
                    for item in input_path.iterdir():
                        if item.is_file() and any(
                            item.name.lower().endswith(ext) for ext in extensions
                        ):
                            video_files.append(item)
                except PermissionError:
                    add_job_log(job_id, f"Permission denied accessing {input_path}", "error")

        _update_job(job_id, total_videos=len(video_files))
        add_job_log(job_id, f"Found {len(video_files)} videos to process", "info")

        # Create output subdirectories
        poses_dir = output_root / "poses"
        features_dir = output_root / "features"
        segments_dir = output_root / "segments"

        poses_dir.mkdir(exist_ok=True)
        features_dir.mkdir(exist_ok=True)
        segments_dir.mkdir(exist_ok=True)

        add_job_log(job_id, "Output directories created", "info")

        # Process each video
        for i, video_path in enumerate(video_files):
            # Check cancellation status thread-safely
            if check_job_cancelled(job_id):
                add_job_log(job_id, "Job cancelled by user", "warning")
                break

            _update_job(job_id, current_video=video_path.name, processed_videos=i)

            add_job_log(job_id, f"Processing [{i+1}/{len(video_files)}]: {video_path.name}", "info")

            try:
                # Get fresh job reference for process_single_video
                current_job = _get_job(job_id)
                if current_job is None:
                    break
                await process_single_video(
                    video_path,
                    output_root,
                    config,
                    current_job,
                    job_id,
                )
                with _jobs_lock:
                    jobs[job_id].successful.append(str(video_path))
                add_job_log(job_id, f"Completed: {video_path.name}", "info")
            except ExtractionCancelled:
                add_job_log(job_id, f"Cancelled during: {video_path.name}", "warning")
                break
            except Exception as e:
                with _jobs_lock:
                    jobs[job_id].failed.append(
                        {
                            "path": str(video_path),
                            "error": str(e),
                        }
                    )
                add_job_log(job_id, f"Failed: {video_path.name} - {str(e)}", "error")

        # Final status update
        with _jobs_lock:
            job = jobs.get(job_id)
            if job:
                job.processed_videos = len(job.successful) + len(job.failed)
                job.status = "completed" if job.status != "cancelled" else "cancelled"
                job.completed_at = datetime.now().isoformat()
                final_status = job.status
                successful_count = len(job.successful)
                failed_count = len(job.failed)

        if final_status == "cancelled":
            add_job_log(
                job_id,
                f"Pipeline cancelled: {successful_count} completed, {failed_count} failed",
                "warning",
            )
        else:
            add_job_log(
                job_id,
                f"Pipeline finished: {successful_count} successful, {failed_count} failed",
                "info",
            )

    except ExtractionCancelled:
        _update_job(job_id, status="cancelled", completed_at=datetime.now().isoformat())
        add_job_log(job_id, "Pipeline cancelled by user", "warning")
        cleanup_gpu_memory()
        add_job_log(job_id, "GPU memory cleanup completed", "info")
    except Exception as e:
        with _jobs_lock:
            if job_id in jobs:
                jobs[job_id].status = "failed"
                jobs[job_id].failed.append({"path": "pipeline", "error": str(e)})
                jobs[job_id].completed_at = datetime.now().isoformat()
        add_job_log(job_id, f"Pipeline error: {str(e)}", "error")
        cleanup_gpu_memory()
        add_job_log(job_id, "GPU memory cleanup completed", "info")


def _run_mediapipe_extraction(
    video_path: Path,
    output_path: Path,
    config: ExtractionConfig,
    job_id: str,
) -> tuple[bool, str | None]:
    """Run MediaPipe extraction in a separate thread. Returns (success, error_message)."""
    try:
        from sltk.extraction.mediapipe import MediaPipeConfig, MediaPipeExtractor

        mp_config = MediaPipeConfig(
            device=config.device,
            enable_hands=config.mediapipe_hands,
            enable_face=config.mediapipe_face,
            compress=config.compress,
        )
        extractor = MediaPipeExtractor(mp_config)
        try:
            extractor.extract_from_video(video_path, output_path)
            add_job_log(job_id, f"  -> MediaPipe output: {output_path}", "info")
            return True, None
        finally:
            extractor.close()
    except Exception as e:
        add_job_log(job_id, f"  -> MediaPipe error: {e}", "error")
        return False, str(e)


def _run_wilor_extraction(
    video_path: Path,
    output_path: Path,
    config: ExtractionConfig,
    job_id: str,
) -> tuple[bool, str | None]:
    """Run WiLoR extraction in a separate thread. Returns (success, error_message)."""
    try:
        from sltk.extraction.wilor import WiLoRConfig, WiLoRExtractor

        wilor_config = WiLoRConfig(
            device=config.device,
            checkpoint_path=config.wilor_checkpoint,
            detector_path=config.wilor_detector,
            compress=config.compress,
            img_batch_size=config.wilor_img_batch_size,
            det_batch_size=config.wilor_det_batch_size,
        )
        extractor = WiLoRExtractor(wilor_config)
        try:
            result = extractor.extract_from_video(video_path, output_path)

            # Check if extraction actually succeeded
            if result.errors:
                error_msg = "; ".join(result.errors)
                raise RuntimeError(f"WiLoR extraction failed: {error_msg}")

            # Verify file was created
            if not output_path.exists():
                raise RuntimeError(f"WiLoR output file not created: {output_path}")

            add_job_log(job_id, f"  -> WiLoR output: {output_path}", "info")
            return True, None
        finally:
            extractor.close()
    except ImportError as e:
        add_job_log(job_id, f"  -> WiLoR not available: {e}", "warning")
        return False, f"Not available: {e}"
    except Exception as e:
        add_job_log(job_id, f"  -> WiLoR error: {e}", "error")
        return False, str(e)


async def process_single_video(
    video_path: Path,
    output_root: Path,
    config: ExtractionConfig,
    job: ExtractionJob,
    job_id: str = "",
):
    """Process a single video through the extraction pipeline."""
    video_stem = video_path.stem

    poses_dir = output_root / "poses"
    features_dir = output_root / "features"
    segments_dir = output_root / "segments"

    # Check for cancellation before starting
    raise_if_cancelled(job_id)

    # Get video info for progress tracking
    job.video_start_time = datetime.now().isoformat()
    job.current_frame = 0
    job.total_frames = 0
    job.stage_errors = {}

    # Determine which stages will run
    active_stages = []
    if config.enable_mediapipe:
        active_stages.append("mediapipe")
    if config.enable_nlf:
        active_stages.append("nlf")
    if config.enable_wilor:
        active_stages.append("wilor")
    active_stages.append("features")  # Always run if prerequisites exist
    if config.enable_segmentation:
        active_stages.append("segmentation")

    job.stages = active_stages
    job.stage_index = 0

    # Initialize stage statuses
    job.stage_statuses = {
        stage: StageStatus(name=stage, status="pending") for stage in active_stages
    }

    try:
        from sltk.extraction.base import get_video_info as get_vid_info

        video_info = get_vid_info(video_path)
        job.total_frames = video_info.get("num_frames", 0)
        add_job_log(
            job_id,
            f"  -> Video info: {job.total_frames} frames at {video_info.get('fps', 25):.1f} fps",
            "info",
        )
    except Exception as e:
        add_job_log(job_id, f"  -> Could not get video info: {e}", "warning")

    # Define output paths
    mediapipe_path = poses_dir / f"{video_stem}_mediapipe.h5"
    wilor_path = poses_dir / f"{video_stem}_wilor.h5"
    nlf_path = poses_dir / f"{video_stem}_nlf.h5"

    # Check which extractors need to run
    run_mediapipe = config.enable_mediapipe and (
        not config.skip_existing or not mediapipe_path.exists()
    )
    run_wilor = config.enable_wilor and (not config.skip_existing or not wilor_path.exists())

    # =========================================================================
    # Sequential extraction: Run MediaPipe then WiLoR to avoid memory issues
    # Running concurrently caused OOM errors due to both loading video frames
    # =========================================================================

    # Stage: MediaPipe extraction
    if run_mediapipe:
        stage_name = "mediapipe"
        job.current_stage = "MediaPipe extraction"
        job.stage_index = job.stages.index(stage_name) if stage_name in job.stages else 0

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].status = "running"
            job.stage_statuses[stage_name].start_time = datetime.now().isoformat()

        add_job_log(job_id, "  -> Running MediaPipe extraction...", "info")
        mp_success, mp_error = _run_mediapipe_extraction(video_path, mediapipe_path, config, job_id)

        if mp_success:
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "completed"
        else:
            job.stage_errors[stage_name] = mp_error or "Unknown error"
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "failed"
                job.stage_statuses[stage_name].error = mp_error

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].end_time = datetime.now().isoformat()

    elif config.enable_mediapipe and not run_mediapipe:
        # MediaPipe was skipped because output exists
        add_job_log(job_id, f"  -> Skipping MediaPipe (exists): {mediapipe_path.name}", "info")
        if "mediapipe" in job.stage_statuses:
            job.stage_statuses["mediapipe"].status = "skipped"

    # Stage: WiLoR extraction
    if run_wilor:
        stage_name = "wilor"
        job.current_stage = "WiLoR extraction"
        job.stage_index = job.stages.index(stage_name) if stage_name in job.stages else 1

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].status = "running"
            job.stage_statuses[stage_name].start_time = datetime.now().isoformat()

        add_job_log(job_id, "  -> Running WiLoR extraction...", "info")
        wilor_success, wilor_error = _run_wilor_extraction(video_path, wilor_path, config, job_id)

        if wilor_success:
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "completed"
        else:
            job.stage_errors[stage_name] = wilor_error or "Unknown error"
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "failed"
                job.stage_statuses[stage_name].error = wilor_error

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].end_time = datetime.now().isoformat()

    elif config.enable_wilor and not run_wilor:
        # WiLoR was skipped because output exists
        add_job_log(job_id, f"  -> Skipping WiLoR (exists): {wilor_path.name}", "info")
        if "wilor" in job.stage_statuses:
            job.stage_statuses["wilor"].status = "skipped"

    # Check for cancellation between stages
    raise_if_cancelled(job_id)

    # Stage: NLF extraction (full-body SMPL-X) - Always runs sequentially (GPU intensive)
    if config.enable_nlf:
        stage_name = "nlf"
        job.current_stage = "NLF extraction"
        job.stage_index = job.stages.index(stage_name) if stage_name in job.stages else 1

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].status = "running"
            job.stage_statuses[stage_name].start_time = datetime.now().isoformat()

        if not config.skip_existing or not nlf_path.exists():
            add_job_log(job_id, "  -> Running NLF extraction...", "info")
            try:
                from sltk.extraction.nlf import NLFConfig, NLFExtractor

                nlf_config = NLFConfig(
                    device=config.device,
                    model_path=config.nlf_checkpoint,
                    smplx_path=config.nlf_smplx_path,
                    compress=config.compress,
                )
                extractor = NLFExtractor(nlf_config)
                try:
                    result = extractor.extract_from_video(video_path, nlf_path)

                    # Check if extraction actually succeeded
                    if result.errors:
                        error_msg = "; ".join(result.errors)
                        raise RuntimeError(f"NLF extraction failed: {error_msg}")

                    # Verify file was created
                    if not nlf_path.exists():
                        raise RuntimeError(f"NLF output file not created: {nlf_path}")

                    add_job_log(job_id, f"  -> NLF output: {nlf_path}", "info")
                    if stage_name in job.stage_statuses:
                        job.stage_statuses[stage_name].status = "completed"
                finally:
                    extractor.close()
            except ImportError as e:
                job.stage_errors[stage_name] = f"Not available: {e}"
                if stage_name in job.stage_statuses:
                    job.stage_statuses[stage_name].status = "failed"
                    job.stage_statuses[stage_name].error = str(e)
                add_job_log(job_id, f"  -> NLF not available: {e}", "warning")
            except Exception as e:
                job.stage_errors[stage_name] = str(e)
                if stage_name in job.stage_statuses:
                    job.stage_statuses[stage_name].status = "failed"
                    job.stage_statuses[stage_name].error = str(e)
                add_job_log(job_id, f"  -> NLF error: {e}", "error")
        else:
            add_job_log(job_id, f"  -> Skipping NLF (exists): {nlf_path.name}", "info")
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "skipped"

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].end_time = datetime.now().isoformat()

    # Check for cancellation between stages
    raise_if_cancelled(job_id)

    # Stage: Feature computation (angles, HaMeR)
    stage_name = "features"
    job.current_stage = "Feature computation"
    job.stage_index = job.stages.index(stage_name) if stage_name in job.stages else 2

    if stage_name in job.stage_statuses:
        job.stage_statuses[stage_name].status = "running"
        job.stage_statuses[stage_name].start_time = datetime.now().isoformat()

    angles_dir = features_dir / "angles"
    hamer_dir = features_dir / "hamer"
    angles_dir.mkdir(exist_ok=True)
    hamer_dir.mkdir(exist_ok=True)

    angles_path = angles_dir / f"{video_stem}.pt"
    hamer_path = hamer_dir / f"{video_stem}.lzma"

    # Check if we have both MediaPipe and WiLoR outputs
    mediapipe_path = poses_dir / f"{video_stem}_mediapipe.h5"
    wilor_path = poses_dir / f"{video_stem}_wilor.h5"

    if mediapipe_path.exists() and wilor_path.exists():
        if not config.skip_existing or not (angles_path.exists() and hamer_path.exists()):
            add_job_log(job_id, "  -> Computing angle and HaMeR features...", "info")
            try:
                from sltk.processing.features import load_features_from_h5, save_features

                angle_features, hamer_features = load_features_from_h5(
                    mediapipe_path=mediapipe_path,
                    wilor_path=wilor_path,
                )

                save_features(
                    angle_features=angle_features,
                    hamer_features=hamer_features,
                    angles_path=angles_path,
                    hamer_path=hamer_path,
                )

                feature_msg = (
                    f"  -> Features: angles={angle_features.shape}, "
                    f"hamer={hamer_features.shape}"
                )
                add_job_log(job_id, feature_msg, "info")
                if stage_name in job.stage_statuses:
                    job.stage_statuses[stage_name].status = "completed"
            except Exception as e:
                job.stage_errors[stage_name] = str(e)
                if stage_name in job.stage_statuses:
                    job.stage_statuses[stage_name].status = "failed"
                    job.stage_statuses[stage_name].error = str(e)
                add_job_log(job_id, f"  -> Feature computation error: {e}", "warning")
        else:
            add_job_log(
                job_id,
                f"  -> Skipping features (exist): {angles_path.name}, {hamer_path.name}",
                "info",
            )
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "skipped"
    else:
        missing = []
        if not mediapipe_path.exists():
            missing.append("MediaPipe")
        if not wilor_path.exists():
            missing.append("WiLoR")
        job.stage_errors[stage_name] = f"Missing prerequisites: {', '.join(missing)}"
        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].status = "skipped"
            job.stage_statuses[stage_name].error = f"Missing: {', '.join(missing)}"
        add_job_log(job_id, f"  -> Skipping features (missing: {', '.join(missing)})", "warning")

    if stage_name in job.stage_statuses:
        job.stage_statuses[stage_name].end_time = datetime.now().isoformat()

    # Check for cancellation between stages
    raise_if_cancelled(job_id)

    # Stage 4: Segmentation (v2 — uses WiLoR H5 directly)
    if config.enable_segmentation:
        stage_name = "segmentation"
        job.current_stage = "Segmentation"
        job.stage_index = job.stages.index(stage_name) if stage_name in job.stages else 3

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].status = "running"
            job.stage_statuses[stage_name].start_time = datetime.now().isoformat()

        segment_path = segments_dir / f"{video_stem}.eaf"

        if wilor_path.exists():
            if not config.skip_existing or not segment_path.exists():
                add_job_log(job_id, "  -> Running sign segmentation (v2)...", "info")
                try:
                    from sltk.segmentation import extract_segments, get_runner
                    from sltk.segmentation.h5_loader import h5_to_features
                    from sltk.segmentation.output import save_as_elan

                    runner = get_runner()
                    features = h5_to_features(wilor_path)
                    labels = runner.predict(features, include_begin=True)

                    # Save as ELAN file
                    save_as_elan({video_stem: labels}, segment_path, fps=25.0)

                    segments = extract_segments(labels)
                    seg_msg = f"  -> Segmented: {len(segments)} segments -> {segment_path.name}"
                    add_job_log(job_id, seg_msg, "info")
                    if stage_name in job.stage_statuses:
                        job.stage_statuses[stage_name].status = "completed"
                except ImportError as e:
                    job.stage_errors[stage_name] = f"Not available: {e}"
                    if stage_name in job.stage_statuses:
                        job.stage_statuses[stage_name].status = "failed"
                        job.stage_statuses[stage_name].error = str(e)
                    add_job_log(job_id, f"  -> Segmentation not available: {e}", "warning")
                except Exception as e:
                    job.stage_errors[stage_name] = str(e)
                    if stage_name in job.stage_statuses:
                        job.stage_statuses[stage_name].status = "failed"
                        job.stage_statuses[stage_name].error = str(e)
                    add_job_log(job_id, f"  -> Segmentation error: {e}", "error")
            else:
                add_job_log(
                    job_id, f"  -> Skipping segmentation (exists): {segment_path.name}", "info"
                )
                if stage_name in job.stage_statuses:
                    job.stage_statuses[stage_name].status = "skipped"
        else:
            error_msg = "WiLoR H5 not available for segmentation"
            job.stage_errors[stage_name] = error_msg
            if stage_name in job.stage_statuses:
                job.stage_statuses[stage_name].status = "skipped"
                job.stage_statuses[stage_name].error = error_msg
            add_job_log(job_id, "  -> Segmentation skipped: WiLoR H5 not available", "warning")

        if stage_name in job.stage_statuses:
            job.stage_statuses[stage_name].end_time = datetime.now().isoformat()


# ============================================================================
# Export for admin router
# ============================================================================


def get_job_stats() -> dict:
    """Get job statistics for admin endpoints."""
    with _jobs_lock:
        total_jobs = len(jobs)
        status_counts: dict[str, int] = {}
        for job in jobs.values():
            status_counts[job.status] = status_counts.get(job.status, 0) + 1

    with _logs_lock:
        total_log_entries = sum(len(logs) for logs in job_logs.values())

    return {
        "total_jobs": total_jobs,
        "status_counts": status_counts,
        "total_log_entries": total_log_entries,
        "config": {
            "ttl_hours": JOB_TTL_HOURS,
            "max_jobs": MAX_JOBS,
            "cleanup_interval_seconds": JOB_CLEANUP_INTERVAL_SECONDS,
        },
    }
