"""
Segmentation v2 API endpoints for SLTK.

Provides endpoints for:
- Model status check
- Single H5 file segmentation (WiLoR features -> BIO labels -> segments)
- Batch directory segmentation with optional output saving
"""

from __future__ import annotations

import asyncio
import logging
from pathlib import Path

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from sltk.api.dependencies import safe_error_detail, validate_path

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


# ============================================================================
# Pydantic Models
# ============================================================================


class SegmentationStatusResponse(BaseModel):
    """Model availability and device info."""

    checkpoint_exists: bool
    device: str | None = None
    gpu_available: bool = False
    error: str | None = None


class SegmentInfo(BaseModel):
    """A single detected sign segment."""

    start_frame: int
    end_frame: int
    start_sec: float
    end_sec: float


class SegmentRequest(BaseModel):
    """Request for single H5 file segmentation."""

    h5_path: str = Field(..., description="Path to WiLoR H5 file")
    include_begin: bool = Field(default=True, description="Include BEGIN (2) labels")
    fps: float = Field(default=25.0, gt=0, le=120, description="FPS for timestamp conversion")


class SegmentResponse(BaseModel):
    """Response from single file segmentation."""

    video_name: str
    num_frames: int
    segments: list[SegmentInfo]
    labels: list[int]


class BatchSegmentRequest(BaseModel):
    """Request for batch directory segmentation."""

    directory: str = Field(..., description="Path to directory with H5 files")
    include_begin: bool = Field(default=True, description="Include BEGIN labels")
    fps: float = Field(default=25.0, gt=0, le=120, description="FPS for timestamp conversion")
    output_path: str | None = Field(default=None, description="Optional path to save results")
    output_format: str = Field(default="json", description="Output format: json, pt, or elan")


class BatchSegmentResponse(BaseModel):
    """Response from batch segmentation."""

    results: dict[str, list[SegmentInfo]]
    total_videos: int


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


@router.get("/status", response_model=SegmentationStatusResponse)
async def get_segmentation_status():
    """
    Check segmentation model availability.

    Returns device info and checkpoint status without triggering model load.
    """
    try:
        import torch

        gpu_available = torch.cuda.is_available()
    except ImportError:
        return SegmentationStatusResponse(
            checkpoint_exists=False,
            gpu_available=False,
            error="PyTorch not installed",
        )

    # Check checkpoint without loading the model
    try:
        from sltk.extraction.weights import get_weight_path

        ckpt_path = get_weight_path("segmentor")
        checkpoint_exists = ckpt_path is not None and Path(ckpt_path).exists()
    except Exception:
        checkpoint_exists = False

    # Report device from runner if already loaded
    from sltk.segmentation.runner import _runner

    device = str(_runner.device) if _runner is not None else None

    return SegmentationStatusResponse(
        checkpoint_exists=checkpoint_exists,
        device=device,
        gpu_available=gpu_available,
    )


@router.post("/segment", response_model=SegmentResponse)
async def segment_single(body: SegmentRequest):
    """
    Run segmentation on a single WiLoR H5 file.

    Loads features from the H5 file, runs BIO prediction, then extracts
    segment boundaries with timestamps.
    """
    h5_path = validate_path(body.h5_path, must_exist=True)

    if not h5_path.suffix.lower() == ".h5":
        raise HTTPException(400, "File must be an H5 file")

    def _run():
        import torch

        from sltk.segmentation.h5_loader import h5_to_features
        from sltk.segmentation.postprocess import extract_segments
        from sltk.segmentation.runner import get_runner

        try:
            features = h5_to_features(h5_path)
        except Exception as e:
            raise HTTPException(
                400, f"Failed to load H5 features: {safe_error_detail('H5 load', e)}"
            )

        try:
            runner = get_runner()
            labels = runner.predict(features, include_begin=body.include_begin)
        except FileNotFoundError as e:
            raise HTTPException(503, f"Segmentation model unavailable: {e}")
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.warning("GPU OOM during segmentation")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise HTTPException(
                    503,
                    "GPU out of memory. Try a shorter video.",
                )
            raise

        raw_segments = extract_segments(labels)
        segments = [
            SegmentInfo(
                start_frame=s,
                end_frame=e,
                start_sec=round(s / body.fps, 4),
                end_sec=round(e / body.fps, 4),
            )
            for s, e in raw_segments
        ]

        return SegmentResponse(
            video_name=h5_path.stem,
            num_frames=len(labels),
            segments=segments,
            labels=labels.tolist(),
        )

    try:
        return await asyncio.to_thread(_run)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Segmentation", e))


@router.post("/segment/batch", response_model=BatchSegmentResponse)
async def segment_batch(body: BatchSegmentRequest):
    """
    Run segmentation on all H5 files in a directory.

    Optionally saves results to disk in json, pt, or elan format.
    """
    dir_path = validate_path(body.directory, must_exist=True)

    if not dir_path.is_dir():
        raise HTTPException(400, "Path must be a directory")

    # Validate output format
    valid_formats = {"json", "pt", "elan"}
    if body.output_format not in valid_formats:
        raise HTTPException(
            400,
            f"Unknown output_format '{body.output_format}'. "
            f"Choose from: {sorted(valid_formats)}",
        )

    # Validate output path if provided
    save_path: Path | None = None
    if body.output_path is not None:
        save_path = validate_path(body.output_path, must_exist=False, allow_creation=True)

    def _run():
        import torch

        from sltk.segmentation.output import OutputFormat, save_predictions
        from sltk.segmentation.postprocess import extract_segments
        from sltk.segmentation.runner import get_runner

        try:
            runner = get_runner()
        except FileNotFoundError as e:
            raise HTTPException(503, f"Segmentation model unavailable: {e}")

        try:
            predictions = runner.run_h5(dir_path, include_begin=body.include_begin)
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.warning("GPU OOM during batch segmentation")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise HTTPException(503, "GPU out of memory during batch segmentation.")
            raise

        # Build response
        results: dict[str, list[SegmentInfo]] = {}
        for name, labels in predictions.items():
            raw_segments = extract_segments(labels)
            results[name] = [
                SegmentInfo(
                    start_frame=s,
                    end_frame=e,
                    start_sec=round(s / body.fps, 4),
                    end_sec=round(e / body.fps, 4),
                )
                for s, e in raw_segments
            ]

        # Save if requested
        if save_path is not None:
            fmt = OutputFormat(body.output_format)
            save_predictions(predictions, save_path, fmt, fps=body.fps)
            logger.info("Saved segmentation results to %s", save_path)

        return BatchSegmentResponse(
            results=results,
            total_videos=len(predictions),
        )

    try:
        return await asyncio.to_thread(_run)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Batch segmentation", e))
