"""Processing router for video processing job submission.

Supports segmentation and spotting jobs using the v2 segmenter
(WiLoR H5 features) and SignRep pipeline (dictionary matching).
Jobs run in the background via asyncio tasks.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, HTTPException
from fastapi.responses import FileResponse

from sltk.api.models import (
    ProcessingJobListResponse,
    ProcessingJobStatus,
    ProcessingSubmitRequest,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/processing", tags=["processing"])

# In-memory job storage
_jobs: dict[str, ProcessingJobStatus] = {}
MAX_JOBS = 100

_SLTK_DIR = Path.home() / ".sltk"

VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".mpg", ".mpeg"}


# ============================================================================
# Helpers
# ============================================================================


def _get_output_dir(workspace: str | None) -> Path:
    """Return workspace outputs directory, creating if needed."""
    if not workspace:
        out = _SLTK_DIR / "outputs"
    else:
        out = _SLTK_DIR / "workspaces" / workspace / "outputs"
    out.mkdir(parents=True, exist_ok=True)
    return out


def _find_h5(video_path: Path) -> Path | None:
    """Find the WiLoR H5 feature file for a video.

    Searches:
      1. {dir}/{stem}_wilor.h5
      2. {dir}/poses/{stem}_wilor.h5
      3. {dir}/{stem}/{stem}_wilor.h5
      4. Any .h5 with stem matching *_wilor* in parent dir
    """
    stem = video_path.stem
    parent = video_path.parent
    candidates = [
        parent / f"{stem}_wilor.h5",
        parent / "poses" / f"{stem}_wilor.h5",
        parent / stem / f"{stem}_wilor.h5",
        parent / stem / "wilor.h5",
    ]
    for c in candidates:
        if c.exists():
            return c
    # Glob fallback
    for h5 in parent.glob(f"{stem}*_wilor.h5"):
        if h5.exists():
            return h5
    return None


def _find_video_for_h5(h5_path: Path) -> Path | None:
    """Find the source video file for an H5 feature file."""
    stem = h5_path.stem.replace("_wilor", "")
    for ext in VIDEO_EXTS:
        candidate = h5_path.parent / f"{stem}{ext}"
        if candidate.exists():
            return candidate
        # Check parent directory (if H5 is in poses/ subdir)
        candidate = h5_path.parent.parent / f"{stem}{ext}"
        if candidate.exists():
            return candidate
    return None


async def _auto_ingest_ai_eaf(eaf_path: Path, workspace: str | None):
    """Ingest an AI-generated ELAN file into the corpus DB."""
    if not workspace:
        return
    try:
        from sltk.api.routers.workspace import _ensure_loaded, _manager

        _ensure_loaded()
        db = _manager.get_corpus_db_by_name(workspace)
        await asyncio.to_thread(db.ingest_eaf, eaf_path, source="model")
        logger.info("Auto-ingested AI EAF into corpus: %s", eaf_path.name)
    except Exception:
        logger.exception("Failed to auto-ingest AI EAF: %s", eaf_path)


# ============================================================================
# Background job runners
# ============================================================================


async def _run_segmentation_job(
    job_id: str,
    video_paths: list[str],
    workspace: str | None = None,
    fps: float = 25.0,
):
    """Run segmentation on WiLoR H5 files in the background."""
    job = _jobs.get(job_id)
    if not job:
        return

    job.status = "running"
    job.started_at = datetime.now(timezone.utc).isoformat()
    output_dir = _get_output_dir(workspace)

    try:
        from sltk.segmentation.h5_loader import h5_to_features
        from sltk.segmentation.output import save_as_elan
        from sltk.segmentation.runner import get_runner

        runner = await asyncio.to_thread(get_runner)

        for i, vpath in enumerate(video_paths):
            job.current_video = Path(vpath).name
            job.progress_percent = (i / len(video_paths)) * 100

            video_path = Path(vpath)
            h5_path = _find_h5(video_path)
            if h5_path is None:
                logger.warning("No WiLoR H5 found for: %s", vpath)
                continue

            features = await asyncio.to_thread(h5_to_features, h5_path)
            labels = await asyncio.to_thread(runner.predict, features)

            # Find video for media descriptor
            media = _find_video_for_h5(h5_path) or video_path
            stem = video_path.stem
            eaf_path = output_dir / f"{stem}_segments.eaf"

            await asyncio.to_thread(
                save_as_elan,
                {stem: labels},
                eaf_path,
                fps,
                media_path=media if media.exists() else None,
            )

            job.output_files.append(str(eaf_path))
            job.processed_videos = i + 1
            job.progress_percent = ((i + 1) / len(video_paths)) * 100

            # Auto-ingest into corpus
            await _auto_ingest_ai_eaf(eaf_path, workspace)

        job.status = "completed"
        job.completed_at = datetime.now(timezone.utc).isoformat()
        job.progress_percent = 100.0

    except Exception as e:
        logger.exception("Segmentation job %s failed", job_id)
        job.status = "failed"
        job.error = str(e)
        job.completed_at = datetime.now(timezone.utc).isoformat()


async def _run_spotting_job(
    job_id: str,
    video_paths: list[str],
    workspace: str | None = None,
    dictionary_dirs: list[str] | None = None,
    top_k: int = 5,
    fps: float = 25.0,
):
    """Run segmentation + spotting pipeline in the background.

    For each video: segment first, then spot glosses against dictionaries.
    """
    job = _jobs.get(job_id)
    if not job:
        return

    job.status = "running"
    job.started_at = datetime.now(timezone.utc).isoformat()
    output_dir = _get_output_dir(workspace)

    if not dictionary_dirs:
        job.status = "failed"
        job.error = "No dictionary directories provided for spotting"
        job.completed_at = datetime.now(timezone.utc).isoformat()
        return

    try:
        from sltk.segmentation.h5_loader import h5_to_features
        from sltk.segmentation.output import save_as_elan, save_spotted_elan
        from sltk.segmentation.postprocess import extract_segments
        from sltk.segmentation.runner import get_runner

        runner = await asyncio.to_thread(get_runner)

        # Lazy-load SignRep pipeline
        from sltk.embedding.pipeline import SignRepPipeline

        pipeline = await asyncio.to_thread(SignRepPipeline)

        # Load dictionaries once
        dictionary = await asyncio.to_thread(pipeline.load_dictionary, dictionary_dirs)

        for i, vpath in enumerate(video_paths):
            video_path = Path(vpath)
            stem = video_path.stem
            job.current_video = video_path.name
            job.progress_percent = (i / len(video_paths)) * 100

            # Step 1: Segment
            h5_path = _find_h5(video_path)
            if h5_path is None:
                logger.warning("No WiLoR H5 found for: %s", vpath)
                continue

            features = await asyncio.to_thread(h5_to_features, h5_path)
            labels = await asyncio.to_thread(runner.predict, features)

            # Save segmentation ELAN
            media = _find_video_for_h5(h5_path) or video_path
            seg_eaf_path = output_dir / f"{stem}_segments.eaf"
            await asyncio.to_thread(
                save_as_elan,
                {stem: labels},
                seg_eaf_path,
                fps,
                media_path=media if media.exists() else None,
            )
            job.output_files.append(str(seg_eaf_path))
            await _auto_ingest_ai_eaf(seg_eaf_path, workspace)

            # Step 2: Extract segment boundaries for spotting
            raw_segments = extract_segments(labels)
            if not raw_segments:
                logger.warning("No segments found for %s, skipping spot", stem)
                job.processed_videos = i + 1
                continue

            # Convert to dicts for pipeline.spot()
            segment_dicts = [
                {
                    "segment_id": idx,
                    "original_video_id": stem,
                    "start_frame": s[0],
                    "end_frame": s[1],
                }
                for idx, s in enumerate(raw_segments)
            ]

            # Step 3: Extract continuous features from video
            continuous = await asyncio.to_thread(
                pipeline.extract_continuous,
                video_path=str(video_path),
                stride=4,
            )

            # Step 4: Spot glosses
            spot_result = await asyncio.to_thread(
                pipeline.spot,
                features=continuous,
                segments=segment_dicts,
                dictionary=dictionary,
                top_k=top_k,
                segment_pooling="max",
            )

            # Step 5: Save spotted ELAN
            spotted_eaf_path = output_dir / f"{stem}_spotted.eaf"
            await asyncio.to_thread(
                save_spotted_elan,
                spot_result,
                spotted_eaf_path,
                fps,
                media_path=media if media.exists() else None,
            )
            job.output_files.append(str(spotted_eaf_path))
            await _auto_ingest_ai_eaf(spotted_eaf_path, workspace)

            job.processed_videos = i + 1
            job.progress_percent = ((i + 1) / len(video_paths)) * 100

        job.status = "completed"
        job.completed_at = datetime.now(timezone.utc).isoformat()
        job.progress_percent = 100.0

    except Exception as e:
        logger.exception("Spotting job %s failed", job_id)
        job.status = "failed"
        job.error = str(e)
        job.completed_at = datetime.now(timezone.utc).isoformat()


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


@router.post("/submit", response_model=ProcessingJobStatus)
async def submit_processing(
    request: ProcessingSubmitRequest,
    background_tasks: BackgroundTasks,
):
    """Submit videos for processing.

    For type="segments", runs segmentation on WiLoR H5 files.
    For type="spots", runs segmentation then gloss spotting.
    For type="both", same as "spots" (segment + spot together).
    Returns a job ID immediately; processing runs in the background.
    """
    if len(_jobs) >= MAX_JOBS:
        completed = [jid for jid, j in _jobs.items() if j.status in ("completed", "failed")]
        for jid in completed[:10]:
            _jobs.pop(jid, None)

    job_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc).isoformat()

    job = ProcessingJobStatus(
        job_id=job_id,
        status="pending",
        created_at=now,
        total_videos=len(request.video_paths),
        processed_videos=0,
        type=request.type,
        model=request.model,
        workspace=request.workspace,
    )
    _jobs[job_id] = job

    if request.type == "segments":
        background_tasks.add_task(
            _run_segmentation_job,
            job_id,
            request.video_paths,
            request.workspace,
            request.fps,
        )
    elif request.type in ("spots", "both"):
        background_tasks.add_task(
            _run_spotting_job,
            job_id,
            request.video_paths,
            request.workspace,
            request.dictionary_dirs,
            request.top_k,
            request.fps,
        )

    return job


@router.get("/status/{job_id}", response_model=ProcessingJobStatus)
async def processing_status(job_id: str):
    """Get the status of a processing job."""
    job = _jobs.get(job_id)
    if job is None:
        raise HTTPException(404, f"Job {job_id} not found")
    return job


@router.get("/jobs", response_model=ProcessingJobListResponse)
async def list_processing_jobs():
    """List all recent processing jobs."""
    jobs = sorted(
        _jobs.values(),
        key=lambda j: j.created_at,
        reverse=True,
    )
    return ProcessingJobListResponse(jobs=jobs, total=len(jobs))


@router.get("/output/{job_id}/{filename}")
async def download_output(job_id: str, filename: str):
    """Download a generated output file from a completed job."""
    job = _jobs.get(job_id)
    if job is None:
        raise HTTPException(404, f"Job {job_id} not found")

    # Find matching output file
    for fpath in job.output_files:
        if Path(fpath).name == filename:
            p = Path(fpath)
            if p.exists():
                return FileResponse(
                    p,
                    filename=p.name,
                    media_type="application/xml",
                )
            raise HTTPException(404, f"File no longer exists: {filename}")

    raise HTTPException(404, f"File {filename} not in job outputs")
