"""
SignRep pipeline API endpoints for SLTK.

Provides endpoints for:
- Dictionary extraction (isolated sign video -> pooled 768-dim feature)
- Batch dictionary extraction via job system
- Continuous feature extraction (full video -> dense sliding-window features)
- Gloss spotting (match segments against dictionary features)
- Listing available dictionary directories

Integrates with the unified job system for GPU-intensive batch tasks.
"""

from __future__ import annotations

import asyncio
import logging
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Literal

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

from sltk.api.dependencies import rate_limit, safe_error_detail, validate_path

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

# ============================================================================
# Constants
# ============================================================================

# Feature cache settings (continuous extraction results)
FEATURE_CACHE_TTL = 1800  # 30 minutes
MAX_FEATURE_CACHE_SIZE = 20  # Max cached continuous results

# Dictionary directories to scan for available dictionaries
DICTIONARY_ROOTS = [
    Path("/vol/research/SignFeaturePool/features2"),
]

# ============================================================================
# Thread-safe pipeline singleton
# ============================================================================

_pipeline = None
_pipeline_lock = threading.Lock()
_pipeline_init_error: str | None = None


def _get_pipeline():
    """
    Thread-safe lazy initialization of SignRepPipeline.

    Uses double-checked locking to avoid repeated initialization attempts.
    Returns the pipeline or raises HTTPException if unavailable.
    """
    global _pipeline, _pipeline_init_error

    # Fast path: already initialized
    if _pipeline is not None:
        return _pipeline

    # Slow path: initialize with lock
    with _pipeline_lock:
        # Double-check after acquiring lock
        if _pipeline is not None:
            return _pipeline

        # Check if we already failed to initialize
        if _pipeline_init_error is not None:
            raise HTTPException(
                503,
                f"SignRep pipeline unavailable: {_pipeline_init_error}",
            )

        try:
            from sltk.embedding.pipeline import SignRepPipeline

            _pipeline = SignRepPipeline()
            logger.info("SignRep pipeline initialized successfully")
            return _pipeline
        except ImportError as e:
            _pipeline_init_error = f"SignRep dependencies not installed: {e}"
            logger.error(_pipeline_init_error)
            raise HTTPException(503, _pipeline_init_error)
        except FileNotFoundError as e:
            _pipeline_init_error = f"SignRep checkpoint not found: {e}"
            logger.error(_pipeline_init_error)
            raise HTTPException(503, _pipeline_init_error)
        except Exception as e:
            _pipeline_init_error = f"Failed to initialize pipeline: {e}"
            logger.error(_pipeline_init_error)
            raise HTTPException(503, _pipeline_init_error)


# ============================================================================
# Feature cache (continuous extraction results, thread-safe)
# ============================================================================

# Maps features_id -> (ContinuousResult, video_path, timestamp)
_feature_cache: dict[str, tuple[Any, str, float]] = {}
_feature_cache_lock = threading.Lock()


def _cache_features(result: Any, video_path: str) -> str:
    """
    Cache a ContinuousResult and return a unique features_id.

    Enforces TTL eviction and size limits to prevent unbounded memory growth.
    """
    features_id = str(uuid.uuid4())[:12]

    with _feature_cache_lock:
        # Evict expired entries first
        now = time.time()
        expired = [k for k, (_, _, ts) in _feature_cache.items() if now - ts > FEATURE_CACHE_TTL]
        for k in expired:
            del _feature_cache[k]

        # Evict oldest if over size limit
        while len(_feature_cache) >= MAX_FEATURE_CACHE_SIZE:
            oldest_key = min(_feature_cache, key=lambda k: _feature_cache[k][2])
            del _feature_cache[oldest_key]

        _feature_cache[features_id] = (result, video_path, now)

    return features_id


def _get_cached_features(features_id: str) -> tuple[Any, str] | None:
    """
    Get cached ContinuousResult by features_id if still valid.

    Returns (ContinuousResult, video_path) or None.
    """
    with _feature_cache_lock:
        if features_id in _feature_cache:
            result, video_path, cached_at = _feature_cache[features_id]
            if time.time() - cached_at < FEATURE_CACHE_TTL:
                return result, video_path
            # Expired
            del _feature_cache[features_id]
    return None


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


class DictionaryExtractRequest(BaseModel):
    """Request for single-video dictionary extraction."""

    video_path: str = Field(..., description="Path to isolated sign video")
    method: str = Field(
        default="middle",
        description="Pooling method: middle, max_pool, mean_pool",
    )
    stride: int = Field(default=1, ge=1, le=16, description="Frame stride between windows")


class DictionaryExtractResponse(BaseModel):
    """Response from dictionary extraction."""

    feature: list[float]
    num_windows: int
    active_probs: list[float]
    fps: float
    total_frames: int
    pooling_method: str
    source_path: str


class BatchDictionaryRequest(BaseModel):
    """Request for batch dictionary extraction via job system."""

    input_dir: str = Field(..., description="Directory of sign videos to process")
    output_dir: str = Field(..., description="Directory to save .npz feature files")
    method: str = Field(default="middle", description="Pooling method: middle, max_pool, mean_pool")
    stride: int = Field(default=1, ge=1, le=16, description="Frame stride")
    resume: bool = Field(
        default=False,
        description="Skip videos with existing output files",
    )


class BatchDictionaryResponse(BaseModel):
    """Response from batch dictionary job submission."""

    job_id: str
    status: str
    message: str
    input_dir: str
    total_videos: int | None = None


class ContinuousExtractRequest(BaseModel):
    """Request for continuous feature extraction."""

    video_path: str = Field(..., description="Path to continuous sign video")
    stride: int = Field(
        default=4,
        ge=1,
        le=32,
        description="Frame stride between windows",
    )


class ContinuousExtractResponse(BaseModel):
    """Response from continuous extraction (features cached server-side)."""

    features_id: str = Field(..., description="ID for cached features, use in /spot")
    num_windows: int
    feature_dim: int
    fps: float
    total_frames: int
    stride: int


class InlineSegment(BaseModel):
    """A segment definition for inline spotting requests."""

    segment_id: int
    original_video_id: str = "unknown"
    start_frame: int
    end_frame: int
    start_timestamp: str = ""
    end_timestamp: str = ""


class SpotRequest(BaseModel):
    """Request for gloss spotting."""

    video_path: str | None = Field(
        default=None,
        description="Path to video (for on-the-fly extraction)",
    )
    features_id: str | None = Field(
        default=None,
        description="ID from /continuous/extract cache",
    )
    dictionary_dirs: list[str] = Field(
        ...,
        min_length=1,
        description="List of dictionary .pt files or .npz directories",
    )
    top_k: int = Field(
        default=20,
        ge=1,
        le=100,
        description="Number of top matches per segment",
    )
    stride: int = Field(
        default=4,
        ge=1,
        le=32,
        description="Frame stride (only used if video_path given)",
    )
    segment_pooling: Literal["max", "mean", "softmax_weighted"] = Field(
        default="max",
        description="How to aggregate similarity across windows per segment",
    )
    segments: list[InlineSegment] | None = Field(
        default=None,
        description="Inline segment definitions (alternative to segments_json)",
    )
    segments_json: str | None = Field(
        default=None,
        description="Path to segments JSON file",
    )
    dict_feature_key: str = Field(
        default="best_latent",
        description="Key to read from .npz dictionary files",
    )


class SpotSegmentResponse(BaseModel):
    """A single segment's spotting matches."""

    segment_id: int
    video_id: str
    start_frame: int
    end_frame: int
    start_ms: int
    end_ms: int
    start_timestamp: str = ""
    end_timestamp: str = ""
    num_windows: int
    top_glosses: list[dict[str, Any]]


class SpotResponse(BaseModel):
    """Response from gloss spotting."""

    segments: list[SpotSegmentResponse]
    metadata: dict[str, Any]


class PipelineStatusResponse(BaseModel):
    """Pipeline availability and status."""

    available: bool
    loaded: bool = False
    device: str | None = None
    checkpoint: str | None = None
    feature_dim: int | None = None
    amp_enabled: bool | None = None
    error: str | None = None
    cached_features_count: int = 0


class DictionaryDirInfo(BaseModel):
    """Information about an available dictionary directory."""

    path: str
    name: str
    dataset: str
    num_files: int
    format: str  # "npz" or "pt"


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


@router.get("/status", response_model=PipelineStatusResponse)
async def get_pipeline_status():
    """
    Get SignRep pipeline availability and configuration.

    Returns device info, checkpoint path, and feature cache stats.
    Does not trigger pipeline initialization if not already loaded.
    """
    global _pipeline_init_error

    # Count cached features
    with _feature_cache_lock:
        cached_count = len(_feature_cache)

    # Check if already initialized without triggering init
    if _pipeline is not None:
        return PipelineStatusResponse(
            available=True,
            loaded=True,
            device=str(_pipeline.device),
            checkpoint=None,  # Don't expose internal paths
            feature_dim=768,
            amp_enabled=_pipeline.amp,
            cached_features_count=cached_count,
        )

    if _pipeline_init_error is not None:
        return PipelineStatusResponse(
            available=False,
            error=_pipeline_init_error,
            cached_features_count=cached_count,
        )

    # Not yet initialized, try to check if deps are importable
    try:
        import importlib

        importlib.import_module("sltk.embedding.pipeline")
        return PipelineStatusResponse(
            available=True,
            loaded=False,
            feature_dim=768,
            cached_features_count=cached_count,
        )
    except ImportError as e:
        return PipelineStatusResponse(
            available=False,
            error=f"Dependencies not available: {e}",
            cached_features_count=cached_count,
        )


@router.post("/dictionary/extract", response_model=DictionaryExtractResponse)
@rate_limit("10/minute")
async def extract_dictionary(request: Request, body: DictionaryExtractRequest):
    """
    Extract a pooled dictionary feature from a single isolated sign video.

    The video is processed with a sliding window, and windows are pooled
    into a single 768-dim representation using the chosen method.

    Rate limited to 10 requests per minute.
    """
    # Validate video path
    video_path = validate_path(body.video_path, must_exist=True)

    # Validate pooling method before GPU work
    from sltk.embedding.pipeline import POOLING_METHODS

    if body.method not in POOLING_METHODS:
        raise HTTPException(
            400,
            f"Unknown pooling method '{body.method}'. "
            f"Choose from: {list(POOLING_METHODS.keys())}",
        )

    def _extract():
        import torch

        pipeline = _get_pipeline()
        try:
            return pipeline.extract_dictionary(
                video_path=str(video_path),
                method=body.method,
                stride=body.stride,
            )
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.warning("GPU OOM during dictionary extraction")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise HTTPException(
                    503,
                    "GPU out of memory. Try reducing video length or " "increasing stride.",
                )
            raise

    try:
        result = await asyncio.to_thread(_extract)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Dictionary extraction", e))

    return DictionaryExtractResponse(
        feature=result.feature.tolist(),
        num_windows=len(result.all_features),
        active_probs=result.active_probs.tolist(),
        fps=result.fps,
        total_frames=result.total_frames,
        pooling_method=result.pooling_method,
        source_path=body.video_path,
    )


@router.post("/dictionary/batch/job", response_model=BatchDictionaryResponse)
@rate_limit("2/minute")
async def start_batch_dictionary_job(
    request: Request,
    body: BatchDictionaryRequest,
    background_tasks: BackgroundTasks,
):
    """
    Start a batch dictionary extraction job via the GPU job system.

    Processes all videos in input_dir and saves .npz features to output_dir.
    Returns a job ID immediately; monitor progress via /api/jobs/{job_id}.

    Rate limited to 2 requests per minute.
    """
    # Validate paths
    input_path = validate_path(body.input_dir, must_exist=True)
    output_path = validate_path(body.output_dir, must_exist=False, allow_creation=True)

    # Count videos for progress estimation
    from sltk.embedding.pipeline import VIDEO_EXTENSIONS

    video_files = [
        p for p in input_path.rglob("*") if p.suffix.lower() in VIDEO_EXTENSIONS and p.is_file()
    ]
    total_videos = len(video_files)

    if total_videos == 0:
        raise HTTPException(400, "No video files found in input directory")

    # Import job system
    try:
        from sltk.api.routers.jobs import (
            JobType,
            check_gpu_availability,
            create_gpu_job,
            is_job_cancelled,
            mark_job_cancelled,
            mark_job_completed,
            mark_job_failed,
            mark_job_started,
            update_job_progress,
        )
    except ImportError:
        raise HTTPException(
            503,
            "Job system not available. Use /dictionary/extract for single " "videos instead.",
        )

    # Check GPU availability
    available, error = check_gpu_availability(JobType.SIGNREP_DICTIONARY)
    if not available:
        raise HTTPException(503, f"Cannot start batch job: {error}")

    # Create job
    job, error = create_gpu_job(
        JobType.SIGNREP_DICTIONARY,
        metadata={
            "input_dir": body.input_dir,
            "output_dir": body.output_dir,
            "method": body.method,
            "stride": body.stride,
            "resume": body.resume,
            "total_videos": total_videos,
        },
    )

    if not job:
        raise HTTPException(503, f"Failed to create job: {error}")

    # Background task
    async def run_batch_job():
        import torch

        try:
            pipeline = _get_pipeline()
            mark_job_started(job.job_id, total_items=total_videos)

            results = {}
            errors = []

            for i, video_file in enumerate(video_files):
                if is_job_cancelled(job.job_id):
                    mark_job_cancelled(job.job_id)
                    return

                update_job_progress(
                    job.job_id,
                    processed_items=i,
                    current_step=f"Processing {video_file.name}",
                )

                # Check resume
                if body.resume:
                    out_file = output_path / video_file.relative_to(input_path).with_suffix(".npz")
                    if out_file.exists():
                        continue

                try:
                    result = pipeline.extract_dictionary(
                        video_path=str(video_file),
                        method=body.method,
                        stride=body.stride,
                    )
                    results[video_file.stem] = result

                    # Save to output
                    out_file = output_path / video_file.relative_to(input_path).with_suffix(".npz")
                    out_file.parent.mkdir(parents=True, exist_ok=True)
                    result.save_npz(out_file)

                except RuntimeError as e:
                    if "out of memory" in str(e).lower():
                        logger.warning("GPU OOM for %s, clearing cache", video_file.name)
                        if torch.cuda.is_available():
                            torch.cuda.empty_cache()
                        errors.append(f"{video_file.name}: GPU OOM")
                    else:
                        errors.append(f"{video_file.name}: {e}")
                except Exception as e:
                    errors.append(f"{video_file.name}: {e}")

            mark_job_completed(
                job.job_id,
                result={
                    "processed": len(results),
                    "errors_count": len(errors),
                    "errors": errors[:20],  # Limit error list
                },
            )

        except Exception as e:
            logger.exception("Batch dictionary job %s failed", job.job_id)
            mark_job_failed(job.job_id, str(e))

    background_tasks.add_task(run_batch_job)

    return BatchDictionaryResponse(
        job_id=job.job_id,
        status="started",
        message="Batch dictionary extraction started",
        input_dir=body.input_dir,
        total_videos=total_videos,
    )


@router.post("/continuous/extract", response_model=ContinuousExtractResponse)
@rate_limit("10/minute")
async def extract_continuous(request: Request, body: ContinuousExtractRequest):
    """
    Extract continuous sliding-window features from a video.

    Features are cached server-side and a features_id is returned.
    Use the features_id in a subsequent /spot request to avoid
    re-extracting features when trying different dictionaries or
    spotting parameters.

    Rate limited to 10 requests per minute.
    """
    video_path = validate_path(body.video_path, must_exist=True)

    def _extract():
        import torch

        pipeline = _get_pipeline()
        try:
            return pipeline.extract_continuous(
                video_path=str(video_path),
                stride=body.stride,
            )
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.warning("GPU OOM during continuous extraction")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise HTTPException(
                    503,
                    "GPU out of memory. Try increasing stride or using a " "shorter video.",
                )
            raise

    try:
        result = await asyncio.to_thread(_extract)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Continuous extraction", e))

    # Cache the result
    features_id = _cache_features(result, body.video_path)

    return ContinuousExtractResponse(
        features_id=features_id,
        num_windows=len(result.features),
        feature_dim=result.features.shape[1],
        fps=result.fps,
        total_frames=result.total_frames,
        stride=result.stride,
    )


@router.post("/spot", response_model=SpotResponse)
@rate_limit("5/minute")
async def spot_glosses(request: Request, body: SpotRequest):
    """
    Spot glosses by matching video features against dictionary entries.

    Provide either:
    - video_path: features will be extracted on-the-fly
    - features_id: reuse cached features from /continuous/extract

    Segments can be provided inline or via a JSON file path.

    Rate limited to 5 requests per minute.
    """
    # Validate: must provide video_path or features_id
    if body.video_path is None and body.features_id is None:
        raise HTTPException(
            400,
            "Must provide either 'video_path' (for on-the-fly extraction) "
            "or 'features_id' (from /continuous/extract cache).",
        )

    # Validate: must provide segments or segments_json
    if body.segments is None and body.segments_json is None:
        raise HTTPException(
            400,
            "Must provide either 'segments' (inline list) or "
            "'segments_json' (path to JSON file).",
        )

    # Validate dictionary dirs
    dict_paths = []
    for d in body.dictionary_dirs:
        p = validate_path(d, must_exist=True)
        dict_paths.append(p)

    # Validate segments_json path if provided
    segments_json_path: Path | None = None
    if body.segments_json is not None:
        segments_json_path = validate_path(body.segments_json, must_exist=True)

    # Resolve features
    continuous_result = None
    if body.features_id is not None:
        cached = _get_cached_features(body.features_id)
        if cached is None:
            raise HTTPException(
                404,
                f"Features '{body.features_id}' not found or expired. "
                "Re-run /continuous/extract.",
            )
        continuous_result, _cached_video_path = cached

    def _spot():
        import torch

        pipeline = _get_pipeline()

        # Load dictionary
        try:
            dictionary = pipeline.load_dictionary(
                [str(p) for p in dict_paths],
                feature_key=body.dict_feature_key,
            )
        except FileNotFoundError as e:
            raise HTTPException(400, f"Dictionary loading failed: {e}")

        # Resolve segments to list[dict]
        if segments_json_path is not None:
            from sltk.embedding.pipeline import load_segments

            all_video_segments = load_segments(str(segments_json_path))
            # Flatten all segments
            flat_segments = []
            for segs in all_video_segments.values():
                flat_segments.extend(segs)
            segments_list = flat_segments
        elif body.segments is not None:
            segments_list = [
                {
                    "segment_id": s.segment_id,
                    "original_video_id": s.original_video_id,
                    "start_frame": s.start_frame,
                    "end_frame": s.end_frame,
                    "start_timestamp": s.start_timestamp,
                    "end_timestamp": s.end_timestamp,
                }
                for s in body.segments
            ]
        else:
            segments_list = []

        if len(segments_list) == 0:
            raise HTTPException(400, "No segments provided or found in JSON")

        # Get or extract features
        nonlocal continuous_result
        if continuous_result is None:
            # Extract on-the-fly from video
            video_path = validate_path(body.video_path, must_exist=True)
            try:
                continuous_result = pipeline.extract_continuous(
                    video_path=str(video_path),
                    stride=body.stride,
                )
            except RuntimeError as e:
                if "out of memory" in str(e).lower():
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    raise HTTPException(503, "GPU out of memory during feature extraction.")
                raise

        # Run spotting
        try:
            result = pipeline.spot(
                features=continuous_result,
                segments=segments_list,
                dictionary=dictionary,
                top_k=body.top_k,
                segment_pooling=body.segment_pooling,
            )
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise HTTPException(503, "GPU out of memory during spotting.")
            raise

        return result

    try:
        result = await asyncio.to_thread(_spot)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Gloss spotting", e))

    # Convert to response models
    response_segments = [
        SpotSegmentResponse(
            segment_id=s.segment_id,
            video_id=s.video_id,
            start_frame=s.start_frame,
            end_frame=s.end_frame,
            start_ms=s.start_ms,
            end_ms=s.end_ms,
            start_timestamp=s.start_timestamp,
            end_timestamp=s.end_timestamp,
            num_windows=s.num_windows,
            top_glosses=s.top_glosses,
        )
        for s in result.segments
    ]

    return SpotResponse(
        segments=response_segments,
        metadata=result.metadata,
    )


@router.get("/dictionaries", response_model=list[DictionaryDirInfo])
async def list_dictionaries():
    """
    List available dictionary directories for spotting.

    Scans known feature directories for dictionary .npz/.pt files.
    """

    def _scan():
        dirs = []
        for root in DICTIONARY_ROOTS:
            if not root.exists():
                continue
            try:
                # Scan for dataset/method subdirectories
                for dataset_dir in sorted(root.iterdir()):
                    if not dataset_dir.is_dir():
                        continue
                    for method_dir in sorted(dataset_dir.iterdir()):
                        if not method_dir.is_dir():
                            continue

                        # Count .npz and .pt files
                        npz_count = len(list(method_dir.glob("*.npz")))
                        pt_count = len(list(method_dir.glob("*.pt")))

                        if npz_count == 0 and pt_count == 0:
                            continue

                        fmt = "pt" if pt_count > 0 else "npz"
                        num_files = pt_count if pt_count > 0 else npz_count

                        dirs.append(
                            DictionaryDirInfo(
                                path=str(method_dir),
                                name=method_dir.name,
                                dataset=dataset_dir.name,
                                num_files=num_files,
                                format=fmt,
                            )
                        )
            except (OSError, PermissionError) as e:
                logger.debug("Cannot scan %s: %s", root, e)

        return dirs

    try:
        return await asyncio.to_thread(_scan)
    except Exception as e:
        raise HTTPException(500, safe_error_detail("Dictionary scan", e))
