"""
Embedding visualization API endpoints for SLTK.

Provides endpoints for:
- Computing SignRep embeddings for datasets
- Visualizing embeddings with PCA/t-SNE/UMAP projections
- Video classification (ready-to-embed, needs-segmentation, needs-features)
- Checking embedding cache status

Integrates with the unified job system for GPU task management.
"""

from __future__ import annotations

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

import numpy as np
from fastapi import APIRouter, BackgroundTasks, HTTPException, Request
from pydantic import BaseModel, Field

from sltk.api.dependencies import get_cached_dataset, rate_limit

router = APIRouter(tags=["embeddings"])
logger = logging.getLogger(__name__)

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

# Maximum samples to prevent O(n^2) memory exhaustion with t-SNE
MAX_SAMPLES_ABSOLUTE = 10000
MAX_SAMPLES_DEFAULT = 500
MAX_SAMPLES_LIMIT = 2000

# Perplexity bounds for t-SNE
TSNE_PERPLEXITY_MIN = 5
TSNE_PERPLEXITY_MAX = 50

# Cache settings
EMBEDDING_CACHE_TTL = 3600  # 1 hour

# ============================================================================
# Thread-safe SignRep initialization
# ============================================================================

_signrep_embedder = None
_signrep_lock = threading.Lock()
_signrep_init_error: str | None = None


def _get_signrep_embedder():
    """
    Thread-safe lazy initialization of SignRep embedder.

    Returns the embedder or raises HTTPException if unavailable.
    """
    global _signrep_embedder, _signrep_init_error

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

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

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

        try:
            from sltk.embedding.signrep import SignRepEmbedder

            _signrep_embedder = SignRepEmbedder(lazy_load=True)
            logger.info("SignRep embedder initialized successfully")
            return _signrep_embedder
        except ImportError as e:
            _signrep_init_error = f"SignRep not installed: {e}"
            logger.error(_signrep_init_error)
            raise HTTPException(503, _signrep_init_error)
        except Exception as e:
            _signrep_init_error = f"Failed to initialize SignRep: {e}"
            logger.error(_signrep_init_error)
            raise HTTPException(503, _signrep_init_error)


# ============================================================================
# Embedding cache (thread-safe)
# ============================================================================

_embedding_cache: dict[str, tuple[np.ndarray, dict[str, Any], float]] = {}
_embedding_cache_lock = threading.Lock()


def _get_cached_embeddings(dataset_name: str) -> tuple[np.ndarray, dict[str, Any]] | None:
    """Get cached embeddings for a dataset if fresh."""
    with _embedding_cache_lock:
        if dataset_name in _embedding_cache:
            embeddings, metadata, cached_at = _embedding_cache[dataset_name]
            if time.time() - cached_at < EMBEDDING_CACHE_TTL:
                return embeddings, metadata
            # Expired, remove from cache
            del _embedding_cache[dataset_name]
    return None


def _cache_embeddings(dataset_name: str, embeddings: np.ndarray, metadata: dict[str, Any]) -> None:
    """Cache embeddings for a dataset."""
    with _embedding_cache_lock:
        _embedding_cache[dataset_name] = (embeddings, metadata, time.time())


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


class EmbeddingPoint(BaseModel):
    """A single point in the embedding visualization."""

    sample_id: str
    x: float
    y: float
    z: float | None = None
    gloss: str | None = None
    signer_id: str | None = None
    video_path: str | None = None
    cluster_id: int = -1
    label: str = ""


class ComputeEmbeddingsRequest(BaseModel):
    """Request to compute embeddings for a dataset."""

    dataset_name: str = Field(..., description="Name of the dataset")
    max_samples: int = Field(
        default=MAX_SAMPLES_DEFAULT,
        ge=1,
        le=MAX_SAMPLES_LIMIT,
        description="Maximum samples to embed",
    )
    aggregation: Literal["mean", "max", "first", "last"] = Field(
        default="mean", description="How to aggregate segment embeddings"
    )
    force_recompute: bool = Field(default=False, description="Force recompute even if cached")


class ComputeEmbeddingsResponse(BaseModel):
    """Response from computing embeddings."""

    dataset_name: str
    num_embeddings: int
    embedding_dim: int
    cached: bool = False
    compute_time_sec: float | None = None
    errors: list[str] = Field(default_factory=list)


class VisualizeEmbeddingsRequest(BaseModel):
    """Request to visualize embeddings."""

    dataset_name: str = Field(..., description="Name of the dataset")
    projection_method: Literal["pca", "tsne", "umap"] = Field(
        default="pca", description="Dimensionality reduction method"
    )
    n_components: Literal[2, 3] = Field(
        default=2, description="Number of dimensions for projection (2 or 3)"
    )
    cluster_method: Literal["hdbscan", "kmeans", "none"] = Field(
        default="none", description="Clustering method"
    )
    n_clusters: int | None = Field(
        default=None, ge=2, le=100, description="Number of clusters (for kmeans)"
    )
    max_samples: int = Field(
        default=MAX_SAMPLES_DEFAULT,
        ge=1,
        le=MAX_SAMPLES_LIMIT,
        description="Maximum samples to visualize",
    )
    # t-SNE specific
    perplexity: int | None = Field(
        default=None,
        ge=TSNE_PERPLEXITY_MIN,
        le=TSNE_PERPLEXITY_MAX,
        description="t-SNE perplexity (auto-adjusted if None)",
    )
    # UMAP specific
    n_neighbors: int = Field(default=15, ge=2, le=100, description="UMAP n_neighbors parameter")
    min_dist: float = Field(default=0.1, ge=0.0, le=1.0, description="UMAP min_dist parameter")


class VisualizeEmbeddingsResponse(BaseModel):
    """Response for embedding visualization."""

    dataset_name: str
    projection_method: str
    num_points: int
    points: list[EmbeddingPoint]
    unique_labels: list[str]
    explained_variance: list[float] | None = None  # For PCA only
    statistics: dict[str, Any] = Field(default_factory=dict)


class EmbeddingStatusResponse(BaseModel):
    """Response for embedding cache status."""

    dataset_name: str
    cached: bool
    num_embeddings: int = 0
    embedding_dim: int = 0
    cached_at: str | None = None
    expires_in_sec: int | None = None
    signrep_available: bool = False
    signrep_loaded: bool = False


# ============================================================================
# Video Classification Models
# ============================================================================


class VideoClassification(str):
    """Video classification for embedding workflow."""

    READY = "ready-to-embed"  # Short dictionary videos with labels
    NEEDS_SEGMENTATION = "needs-segmentation"  # Long videos needing boundary detection
    NEEDS_FEATURES = "needs-features"  # Videos missing pose features
    UNANNOTATED = "unannotated"  # Videos with no annotations


class SampleClassification(BaseModel):
    """Classification result for a single sample."""

    sample_id: str
    classification: str
    video_path: str | None = None
    video_duration_sec: float | None = None
    has_video: bool = False
    has_poses: bool = False
    has_segments: bool = False
    has_glosses: bool = False
    num_segments: int = 0
    reason: str = ""


class DatasetClassificationResponse(BaseModel):
    """Classification results for a dataset."""

    dataset_name: str
    total_samples: int
    ready_count: int = 0
    needs_segmentation_count: int = 0
    needs_features_count: int = 0
    unannotated_count: int = 0
    samples: list[SampleClassification] = Field(default_factory=list)
    prerequisites: dict[str, bool] = Field(default_factory=dict)
    recommended_action: str = ""


class DatasetPrerequisites(BaseModel):
    """Prerequisite status for a dataset."""

    dataset_name: str
    has_videos: bool = False
    has_poses: bool = False
    has_segments: bool = False
    has_embeddings: bool = False
    video_count: int = 0
    pose_count: int = 0
    segment_count: int = 0
    embedding_count: int = 0
    all_ready: bool = False
    next_step: str = ""
    next_step_description: str = ""


class EmbeddingJobRequest(BaseModel):
    """Request to start an embedding computation job."""

    dataset_name: str = Field(..., description="Name of the dataset")
    max_samples: int = Field(
        default=MAX_SAMPLES_DEFAULT,
        ge=1,
        le=MAX_SAMPLES_LIMIT,
        description="Maximum samples to embed",
    )
    aggregation: Literal["mean", "max", "first", "last"] = Field(
        default="mean", description="How to aggregate segment embeddings"
    )


class EmbeddingJobResponse(BaseModel):
    """Response from starting an embedding job."""

    job_id: str | None = None
    status: str
    message: str
    cached: bool = False
    num_embeddings: int = 0


# ============================================================================
# Helper functions
# ============================================================================


def _compute_embeddings_sync(
    dataset_name: str,
    max_samples: int,
    aggregation: str,
) -> tuple[np.ndarray, dict[str, Any], list[str]]:
    """
    Synchronous embedding computation (runs in thread pool).

    Returns:
        Tuple of (embeddings array, metadata dict, error list)
    """
    embedder = _get_signrep_embedder()

    import torch

    dataset = get_cached_dataset(dataset_name, split="all")

    embeddings_list = []
    sample_ids = []
    glosses = []
    signer_ids = []
    video_paths = []
    errors = []

    # Get sample indices
    if hasattr(dataset, "_index"):
        all_sample_ids = dataset._index[:max_samples]
    else:
        all_sample_ids = list(range(min(len(dataset), max_samples)))

    for sample_id in all_sample_ids:
        try:
            sample = dataset._load_sample(sample_id)

            # Skip if no video
            if sample.video_path is None or not Path(sample.video_path).exists():
                continue

            # Compute embedding with GPU OOM handling
            try:
                embedding = embedder.embed_video(sample.video_path, aggregation=aggregation)
            except RuntimeError as e:
                if "out of memory" in str(e).lower() or "CUDA" in str(e):
                    # GPU OOM - try to recover and fall back to CPU
                    logger.warning(f"GPU OOM for {sample_id}, attempting CPU fallback")
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

                    # Try with CPU (if embedder supports device switching)
                    try:
                        # Create a new config with CPU
                        from sltk.embedding.signrep import SignRepConfig, SignRepEmbedder

                        cpu_config = SignRepConfig(device="cpu")
                        cpu_embedder = SignRepEmbedder(config=cpu_config, lazy_load=False)
                        embedding = cpu_embedder.embed_video(
                            sample.video_path, aggregation=aggregation
                        )
                        logger.info(f"Successfully computed embedding on CPU for {sample_id}")
                    except Exception as cpu_e:
                        errors.append(f"{sample_id}: CPU fallback failed - {cpu_e}")
                        continue
                else:
                    raise

            embeddings_list.append(embedding)
            sample_ids.append(str(sample.id))

            # Collect metadata
            if sample.glosses:
                glosses.append(sample.glosses[0])
            else:
                glosses.append("")

            signer_ids.append(sample.signer_id or "")
            video_paths.append(str(sample.video_path) if sample.video_path else "")

        except Exception as e:
            errors.append(f"{sample_id}: {str(e)}")
            logger.debug(f"Failed to embed sample {sample_id}: {e}")

    if not embeddings_list:
        raise HTTPException(400, f"No valid embeddings computed for {dataset_name}")

    embeddings = np.stack(embeddings_list)
    metadata = {
        "sample_ids": sample_ids,
        "glosses": glosses,
        "signer_ids": signer_ids,
        "video_paths": video_paths,
    }

    return embeddings, metadata, errors


def _reduce_dimensions_sync(
    embeddings: np.ndarray, method: str, n_components: int, **kwargs
) -> tuple[np.ndarray, list[float] | None]:
    """
    Synchronous dimensionality reduction (runs in thread pool).

    Returns:
        Tuple of (projections, explained_variance or None)
    """
    n_samples = embeddings.shape[0]
    explained_variance = None

    if method == "pca":
        try:
            from sklearn.decomposition import PCA
        except ImportError:
            raise HTTPException(500, "scikit-learn required for PCA: pip install scikit-learn")

        pca = PCA(n_components=n_components)
        projections = pca.fit_transform(embeddings)
        # Handle potential NaN values in explained_variance_ratio_ (can occur with small datasets)
        variance_ratio = pca.explained_variance_ratio_
        if np.any(np.isnan(variance_ratio)):
            logger.warning("PCA produced NaN variance ratios, using zeros")
            variance_ratio = np.nan_to_num(variance_ratio, nan=0.0)
        explained_variance = variance_ratio.tolist()

    elif method == "tsne":
        try:
            from sklearn.manifold import TSNE
        except ImportError:
            raise HTTPException(500, "scikit-learn required for t-SNE: pip install scikit-learn")

        # Validate perplexity
        perplexity = kwargs.get("perplexity", 30)
        # Perplexity must be less than n_samples
        max_perplexity = max(TSNE_PERPLEXITY_MIN, min(n_samples - 1, TSNE_PERPLEXITY_MAX))
        if perplexity >= n_samples:
            perplexity = max_perplexity
            logger.warning(
                f"Adjusted t-SNE perplexity to {perplexity} (must be < n_samples={n_samples})"
            )

        # Safety check: t-SNE is O(n^2) memory
        if n_samples > MAX_SAMPLES_ABSOLUTE:
            raise HTTPException(
                400,
                f"t-SNE limited to {MAX_SAMPLES_ABSOLUTE} samples to prevent memory exhaustion. "
                f"Got {n_samples} samples. Use PCA or UMAP for larger datasets.",
            )

        tsne = TSNE(
            n_components=n_components,
            perplexity=perplexity,
            max_iter=kwargs.get("n_iter", 1000),
            random_state=42,
        )
        projections = tsne.fit_transform(embeddings)

    elif method == "umap":
        try:
            import umap
        except ImportError:
            raise HTTPException(500, "umap-learn required for UMAP: pip install umap-learn")

        reducer = umap.UMAP(
            n_components=n_components,
            n_neighbors=kwargs.get("n_neighbors", 15),
            min_dist=kwargs.get("min_dist", 0.1),
            metric="cosine",
            random_state=42,
        )
        projections = reducer.fit_transform(embeddings)

    else:
        raise HTTPException(400, f"Unknown projection method: {method}")

    # Handle any NaN/Inf values in projections (can happen with degenerate data)
    if np.any(~np.isfinite(projections)):
        logger.warning("Dimensionality reduction produced non-finite values, cleaning up")
        projections = np.nan_to_num(projections, nan=0.0, posinf=0.0, neginf=0.0)

    return projections, explained_variance


def _cluster_sync(
    embeddings: np.ndarray,
    method: str,
    n_clusters: int | None = None,
) -> np.ndarray:
    """
    Synchronous clustering (runs in thread pool).

    Returns:
        Array of cluster IDs (-1 for noise points).
    """
    if method == "none":
        return np.full(len(embeddings), -1, dtype=np.int32)

    if method == "kmeans":
        if n_clusters is None:
            raise HTTPException(400, "n_clusters required for kmeans clustering")

        try:
            from sklearn.cluster import KMeans
        except ImportError:
            raise HTTPException(500, "scikit-learn required: pip install scikit-learn")

        # Ensure n_clusters doesn't exceed n_samples
        n_clusters = min(n_clusters, len(embeddings))

        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        return kmeans.fit_predict(embeddings)

    elif method == "hdbscan":
        try:
            import hdbscan
        except ImportError:
            raise HTTPException(500, "hdbscan required: pip install hdbscan")

        clusterer = hdbscan.HDBSCAN(min_cluster_size=5)
        return clusterer.fit_predict(embeddings)

    else:
        raise HTTPException(400, f"Unknown clustering method: {method}")


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


@router.post("/api/embeddings/compute", response_model=ComputeEmbeddingsResponse)
@rate_limit("10/minute")
async def compute_embeddings(request: Request, body: ComputeEmbeddingsRequest):
    """
    Compute SignRep embeddings for a dataset.

    Embeds videos using the SignRep model and caches the results.

    Rate limited to 10 requests per minute (embedding computation is expensive).
    """
    start_time = time.time()

    # Check cache first (unless force recompute)
    if not body.force_recompute:
        cached = _get_cached_embeddings(body.dataset_name)
        if cached is not None:
            embeddings, metadata = cached
            return ComputeEmbeddingsResponse(
                dataset_name=body.dataset_name,
                num_embeddings=len(embeddings),
                embedding_dim=embeddings.shape[1] if len(embeddings) > 0 else 0,
                cached=True,
            )

    # Compute embeddings in thread pool
    embeddings, metadata, errors = await asyncio.to_thread(
        _compute_embeddings_sync,
        body.dataset_name,
        body.max_samples,
        body.aggregation,
    )

    # Cache results
    _cache_embeddings(body.dataset_name, embeddings, metadata)

    compute_time = time.time() - start_time

    return ComputeEmbeddingsResponse(
        dataset_name=body.dataset_name,
        num_embeddings=len(embeddings),
        embedding_dim=embeddings.shape[1],
        cached=False,
        compute_time_sec=round(compute_time, 2),
        errors=errors[:10],  # Limit error list
    )


@router.post("/api/embeddings/visualize", response_model=VisualizeEmbeddingsResponse)
@rate_limit("30/minute")
async def visualize_embeddings(request: Request, body: VisualizeEmbeddingsRequest):
    """
    Get PCA/t-SNE/UMAP projections of embeddings for visualization.

    Returns 2D or 3D coordinates for each sample along with metadata
    for interactive visualization in the frontend.

    Rate limited to 30 requests per minute.
    """
    # Check cache for embeddings
    cached = _get_cached_embeddings(body.dataset_name)
    if cached is None:
        raise HTTPException(
            400,
            f"No embeddings cached for {body.dataset_name}. "
            "Call POST /api/embeddings/compute first.",
        )

    embeddings, metadata = cached
    n_samples = len(embeddings)

    # Apply max_samples limit
    if n_samples > body.max_samples:
        # Randomly sample to respect limit
        np.random.seed(42)  # Reproducible sampling
        indices = np.random.choice(n_samples, body.max_samples, replace=False)
        indices = np.sort(indices)

        embeddings = embeddings[indices]
        metadata = {
            "sample_ids": [metadata["sample_ids"][i] for i in indices],
            "glosses": [metadata["glosses"][i] for i in indices],
            "signer_ids": [metadata["signer_ids"][i] for i in indices],
            "video_paths": [metadata["video_paths"][i] for i in indices],
        }
        n_samples = len(embeddings)

    # Validate t-SNE perplexity
    perplexity = body.perplexity
    if body.projection_method == "tsne":
        if perplexity is None:
            perplexity = min(30, max(TSNE_PERPLEXITY_MIN, n_samples - 1))
        elif perplexity >= n_samples:
            perplexity = max(TSNE_PERPLEXITY_MIN, n_samples - 1)
            logger.warning(f"Adjusted perplexity to {perplexity} (must be < n_samples)")

    # Dimensionality reduction
    projections, explained_variance = await asyncio.to_thread(
        _reduce_dimensions_sync,
        embeddings,
        body.projection_method,
        body.n_components,
        perplexity=perplexity,
        n_neighbors=body.n_neighbors,
        min_dist=body.min_dist,
    )

    # Clustering
    cluster_ids = await asyncio.to_thread(
        _cluster_sync,
        embeddings,
        body.cluster_method,
        body.n_clusters,
    )

    # Build points list (convert numpy to native Python types for JSON serialization)
    points = []
    unique_labels = set()

    for i in range(n_samples):
        gloss = metadata["glosses"][i]
        label = gloss if gloss else f"cluster_{cluster_ids[i]}"
        unique_labels.add(label)

        point = EmbeddingPoint(
            sample_id=metadata["sample_ids"][i],
            x=float(projections[i, 0]),
            y=float(projections[i, 1]),
            z=float(projections[i, 2]) if body.n_components == 3 else None,
            gloss=gloss if gloss else None,
            signer_id=metadata["signer_ids"][i] if metadata["signer_ids"][i] else None,
            video_path=metadata["video_paths"][i] if metadata["video_paths"][i] else None,
            cluster_id=int(cluster_ids[i]),
            label=label,
        )
        points.append(point)

    # Compute statistics
    statistics = {
        "total_samples": n_samples,
        "unique_glosses": len(set(g for g in metadata["glosses"] if g)),
        "unique_signers": len(set(s for s in metadata["signer_ids"] if s)),
    }

    if body.cluster_method != "none":
        unique_clusters = set(cluster_ids)
        num_noise = np.sum(cluster_ids == -1)
        statistics["num_clusters"] = len(unique_clusters) - (1 if -1 in unique_clusters else 0)
        statistics["num_noise_points"] = int(num_noise)

    return VisualizeEmbeddingsResponse(
        dataset_name=body.dataset_name,
        projection_method=body.projection_method,
        num_points=n_samples,
        points=points,
        unique_labels=sorted(list(unique_labels)),
        explained_variance=explained_variance,
        statistics=statistics,
    )


@router.get("/api/embeddings/status/{dataset_name}", response_model=EmbeddingStatusResponse)
async def get_embedding_status(dataset_name: str):
    """
    Check embedding cache status for a dataset.

    Returns information about whether embeddings are cached,
    cache expiration, and SignRep availability.
    """
    # Check SignRep availability
    signrep_available = False
    signrep_loaded = False

    try:
        embedder = _get_signrep_embedder()
        signrep_available = True
        signrep_loaded = embedder.is_loaded()
    except HTTPException:
        pass

    # Check cache
    cached = False
    num_embeddings = 0
    embedding_dim = 0
    cached_at_str = None
    expires_in = None

    with _embedding_cache_lock:
        if dataset_name in _embedding_cache:
            embeddings, metadata, cached_at = _embedding_cache[dataset_name]
            age = time.time() - cached_at

            if age < EMBEDDING_CACHE_TTL:
                cached = True
                num_embeddings = len(embeddings)
                embedding_dim = embeddings.shape[1] if len(embeddings) > 0 else 0

                # Format cached_at as ISO string
                from datetime import datetime

                cached_at_str = datetime.fromtimestamp(cached_at).isoformat()
                expires_in = int(EMBEDDING_CACHE_TTL - age)

    return EmbeddingStatusResponse(
        dataset_name=dataset_name,
        cached=cached,
        num_embeddings=num_embeddings,
        embedding_dim=embedding_dim,
        cached_at=cached_at_str,
        expires_in_sec=expires_in,
        signrep_available=signrep_available,
        signrep_loaded=signrep_loaded,
    )


@router.delete("/api/embeddings/cache/{dataset_name}")
async def clear_embedding_cache(dataset_name: str):
    """
    Clear cached embeddings for a dataset.

    Useful for forcing recomputation after dataset changes.
    """
    with _embedding_cache_lock:
        if dataset_name in _embedding_cache:
            del _embedding_cache[dataset_name]
            return {"status": "cleared", "dataset_name": dataset_name}
        else:
            return {"status": "not_cached", "dataset_name": dataset_name}


@router.delete("/api/embeddings/cache")
async def clear_all_embedding_cache():
    """
    Clear all cached embeddings.

    Useful for freeing memory or after major updates.
    """
    with _embedding_cache_lock:
        count = len(_embedding_cache)
        _embedding_cache.clear()

    return {"status": "cleared", "datasets_cleared": count}


@router.get("/api/embeddings/signrep/status")
async def get_signrep_status():
    """
    Get SignRep model status and configuration.

    Returns information about whether SignRep is available,
    loaded, and configuration details.
    """
    global _signrep_init_error

    status = {
        "available": False,
        "loaded": False,
        "error": _signrep_init_error,
        "embedding_dim": None,
        "checkpoint_path": None,
    }

    try:
        embedder = _get_signrep_embedder()
        status["available"] = True
        status["loaded"] = embedder.is_loaded()
        status["embedding_dim"] = embedder.embedding_dim
        status["checkpoint_path"] = str(embedder.config.checkpoint_path)
        status["error"] = None
    except HTTPException as e:
        status["error"] = e.detail

    return status


# ============================================================================
# Video Classification Endpoints
# ============================================================================

# Video length thresholds (in seconds)
SHORT_VIDEO_THRESHOLD = 10.0  # Videos under this are "dictionary" style
LONG_VIDEO_THRESHOLD = 60.0  # Videos over this definitely need segmentation


def _classify_sample(
    sample,
    pose_paths: dict[str, Path],
    segment_paths: dict[str, Path],
) -> SampleClassification:
    """
    Classify a single sample for embedding workflow.

    Returns classification based on:
    - Video existence and duration
    - Available pose features
    - Existing segments/annotations
    """
    sample_id = str(sample.id)
    video_path = sample.video_path
    has_video = video_path is not None and Path(video_path).exists()

    # Get video duration
    video_duration = None
    if has_video:
        try:
            from sltk.extraction.base import get_video_info

            info = get_video_info(video_path)
            video_duration = info.get("duration_sec") or info.get("num_frames", 0) / info.get(
                "fps", 25
            )
        except Exception:
            pass

    # Check for poses
    has_poses = sample_id in pose_paths

    # Check for segments
    has_segments = sample_id in segment_paths or (
        hasattr(sample, "segments") and sample.segments and len(sample.segments) > 0
    )

    # Check for glosses
    has_glosses = sample.glosses is not None and len(sample.glosses) > 0
    num_segments = len(sample.segments) if hasattr(sample, "segments") and sample.segments else 0

    # Determine classification
    if not has_video:
        classification = "needs-features"
        reason = "No video file found"
    elif not has_poses:
        classification = "needs-features"
        reason = "Pose features not extracted"
    elif video_duration and video_duration > LONG_VIDEO_THRESHOLD:
        if has_segments or has_glosses:
            classification = "ready-to-embed"
            reason = f"Long video ({video_duration:.1f}s) with segments/annotations"
        else:
            classification = "needs-segmentation"
            reason = f"Long video ({video_duration:.1f}s) needs sign boundaries"
    elif video_duration and video_duration > SHORT_VIDEO_THRESHOLD:
        if has_segments or has_glosses:
            classification = "ready-to-embed"
            reason = "Medium video with segments/annotations"
        else:
            classification = "needs-segmentation"
            reason = "Video may contain multiple signs, segmentation recommended"
    else:
        # Short video - likely dictionary entry
        if has_glosses:
            classification = "ready-to-embed"
            reason = "Short video with gloss label"
        elif has_segments:
            classification = "ready-to-embed"
            reason = "Short video with segments"
        else:
            classification = "unannotated"
            reason = "No annotations available (will embed whole video)"

    return SampleClassification(
        sample_id=sample_id,
        classification=classification,
        video_path=str(video_path) if video_path else None,
        video_duration_sec=video_duration,
        has_video=has_video,
        has_poses=has_poses,
        has_segments=has_segments,
        has_glosses=has_glosses,
        num_segments=num_segments,
        reason=reason,
    )


def _find_feature_files(dataset_name: str, feature_type: str) -> dict[str, Path]:
    """Find feature files for a dataset."""
    result = {}

    # Check common feature locations
    feature_roots = [
        Path(f"/vol/research/SignFeaturePool/features2/{dataset_name.lower()}/{feature_type}"),
        Path(f"/vol/research/SignFeaturePool/features2/{dataset_name}/{feature_type}"),
    ]

    for root in feature_roots:
        if root.exists():
            for f in root.glob("**/*"):
                if f.is_file() and f.suffix in (".h5", ".pt", ".lzma"):
                    sample_id = f.stem.replace(f"_{feature_type}", "")
                    result[sample_id] = f

    return result


@router.post("/api/embeddings/classify", response_model=DatasetClassificationResponse)
@rate_limit("30/minute")
async def classify_dataset_for_embedding(
    request: Request,
    dataset_name: str,
    max_samples: int = 100,
):
    """
    Classify dataset samples for embedding workflow.

    Determines which samples are:
    - ready-to-embed: Have poses and appropriate length/annotations
    - needs-segmentation: Long videos without sign boundaries
    - needs-features: Missing pose features
    - unannotated: No labels (can still embed but no semantic grouping)

    This helps the frontend show appropriate UI guidance.
    """
    dataset = get_cached_dataset(dataset_name, split="all")

    # Find feature files
    pose_paths = {}
    pose_paths.update(_find_feature_files(dataset_name, "wilor"))
    pose_paths.update(_find_feature_files(dataset_name, "mediapipe"))

    segment_paths = _find_feature_files(dataset_name, "segments")

    # Get sample IDs
    if hasattr(dataset, "_index"):
        sample_ids = dataset._index[:max_samples]
    else:
        sample_ids = list(range(min(len(dataset), max_samples)))

    # Classify samples
    classifications = []
    ready_count = 0
    needs_seg_count = 0
    needs_feat_count = 0
    unannotated_count = 0

    for sample_id in sample_ids:
        try:
            sample = dataset._load_sample(sample_id)
            classification = await asyncio.to_thread(
                _classify_sample, sample, pose_paths, segment_paths
            )
            classifications.append(classification)

            if classification.classification == "ready-to-embed":
                ready_count += 1
            elif classification.classification == "needs-segmentation":
                needs_seg_count += 1
            elif classification.classification == "needs-features":
                needs_feat_count += 1
            else:
                unannotated_count += 1

        except Exception as e:
            logger.debug(f"Failed to classify {sample_id}: {e}")

    # Determine recommended action
    if needs_feat_count > ready_count:
        recommended_action = "extract_features"
    elif needs_seg_count > ready_count:
        recommended_action = "run_segmentation"
    elif ready_count > 0:
        recommended_action = "compute_embeddings"
    else:
        recommended_action = "add_annotations"

    return DatasetClassificationResponse(
        dataset_name=dataset_name,
        total_samples=len(classifications),
        ready_count=ready_count,
        needs_segmentation_count=needs_seg_count,
        needs_features_count=needs_feat_count,
        unannotated_count=unannotated_count,
        samples=classifications,
        prerequisites={
            "has_videos": ready_count + needs_seg_count > 0,
            "has_poses": ready_count + needs_seg_count > needs_feat_count,
            "has_annotations": ready_count > unannotated_count,
        },
        recommended_action=recommended_action,
    )


@router.get("/api/embeddings/prerequisites/{dataset_name}", response_model=DatasetPrerequisites)
async def get_dataset_prerequisites(dataset_name: str):
    """
    Get prerequisite status for embedding visualization.

    Returns what's available and what's needed before embedding.
    """
    dataset = get_cached_dataset(dataset_name, split="all")

    # Count samples
    total = len(dataset) if hasattr(dataset, "__len__") else 0

    # Check for pose files
    pose_paths = {}
    pose_paths.update(_find_feature_files(dataset_name, "wilor"))
    pose_paths.update(_find_feature_files(dataset_name, "mediapipe"))
    pose_count = len(pose_paths)

    # Check for segment files
    segment_paths = _find_feature_files(dataset_name, "segments")
    segment_count = len(segment_paths)

    # Check for cached embeddings
    cached = _get_cached_embeddings(dataset_name)
    embedding_count = len(cached[0]) if cached else 0

    # Check for videos by sampling
    video_count = 0
    sample_limit = min(100, total)
    sample_ids = (
        dataset._index[:sample_limit] if hasattr(dataset, "_index") else list(range(sample_limit))
    )

    for sample_id in sample_ids:
        try:
            sample = dataset._load_sample(sample_id)
            if sample.video_path and Path(sample.video_path).exists():
                video_count += 1
        except Exception:
            pass

    # Scale video count to estimate total
    if sample_limit > 0:
        video_count = int(video_count * (total / sample_limit))

    # Determine next step
    has_videos = video_count > 0
    has_poses = pose_count > 0
    has_segments = segment_count > 0
    has_embeddings = embedding_count > 0

    if not has_videos:
        next_step = "connect_videos"
        next_step_description = "Connect video files to the dataset"
    elif not has_poses:
        next_step = "extract_poses"
        next_step_description = "Extract pose features from videos using MediaPipe or WiLoR"
    elif not has_embeddings:
        next_step = "compute_embeddings"
        next_step_description = "Compute SignRep embeddings for visualization"
    else:
        next_step = "visualize"
        next_step_description = "Ready to visualize embeddings!"

    all_ready = has_videos and has_poses and has_embeddings

    return DatasetPrerequisites(
        dataset_name=dataset_name,
        has_videos=has_videos,
        has_poses=has_poses,
        has_segments=has_segments,
        has_embeddings=has_embeddings,
        video_count=video_count,
        pose_count=pose_count,
        segment_count=segment_count,
        embedding_count=embedding_count,
        all_ready=all_ready,
        next_step=next_step,
        next_step_description=next_step_description,
    )


# ============================================================================
# Job-based Embedding Computation
# ============================================================================


@router.post("/api/embeddings/compute/job", response_model=EmbeddingJobResponse)
@rate_limit("5/minute")
async def start_embedding_job(
    request: Request,
    body: EmbeddingJobRequest,
    background_tasks: BackgroundTasks,
):
    """
    Start an embedding computation job with GPU management.

    Uses the unified job system to:
    - Check GPU availability
    - Prevent concurrent GPU tasks
    - Track progress
    - Support cancellation

    Rate limited to 5 requests per minute.
    """
    # Check cache first
    cached = _get_cached_embeddings(body.dataset_name)
    if cached is not None:
        embeddings, metadata = cached
        return EmbeddingJobResponse(
            job_id=None,
            status="cached",
            message="Embeddings already cached",
            cached=True,
            num_embeddings=len(embeddings),
        )

    # Import job system
    try:
        from sltk.api.routers.jobs import (
            JobCancelledException,
            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:
        # Fallback if jobs router not available
        return EmbeddingJobResponse(
            job_id=None,
            status="error",
            message="Job system not available. Use /api/embeddings/compute instead.",
        )

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

    # Create job
    job, error = create_gpu_job(
        JobType.EMBEDDING,
        metadata={
            "dataset_name": body.dataset_name,
            "max_samples": body.max_samples,
            "aggregation": body.aggregation,
        },
    )

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

    # Start background task
    async def run_embedding_job():
        import torch

        try:
            embedder = _get_signrep_embedder()
            dataset = get_cached_dataset(body.dataset_name, split="all")

            # Get sample IDs
            if hasattr(dataset, "_index"):
                all_sample_ids = dataset._index[: body.max_samples]
            else:
                all_sample_ids = list(range(min(len(dataset), body.max_samples)))

            mark_job_started(job.job_id, total_items=len(all_sample_ids))

            embeddings_list = []
            sample_ids = []
            glosses = []
            signer_ids = []
            video_paths = []
            errors = []

            for i, sample_id in enumerate(all_sample_ids):
                # Check cancellation
                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"Embedding sample {sample_id}",
                )

                try:
                    sample = dataset._load_sample(sample_id)

                    if sample.video_path is None or not Path(sample.video_path).exists():
                        continue

                    # Compute embedding with OOM handling
                    try:
                        embedding = embedder.embed_video(
                            sample.video_path,
                            aggregation=body.aggregation,
                        )
                    except RuntimeError as e:
                        if "out of memory" in str(e).lower():
                            logger.warning(f"GPU OOM for {sample_id}, trying CPU")
                            if torch.cuda.is_available():
                                torch.cuda.empty_cache()

                            from sltk.embedding.signrep import SignRepConfig, SignRepEmbedder

                            cpu_embedder = SignRepEmbedder(
                                config=SignRepConfig(device="cpu"),
                                lazy_load=False,
                            )
                            embedding = cpu_embedder.embed_video(
                                sample.video_path,
                                aggregation=body.aggregation,
                            )
                        else:
                            raise

                    embeddings_list.append(embedding)
                    sample_ids.append(str(sample.id))
                    glosses.append(sample.glosses[0] if sample.glosses else "")
                    signer_ids.append(sample.signer_id or "")
                    video_paths.append(str(sample.video_path) if sample.video_path else "")

                except Exception as e:
                    errors.append(f"{sample_id}: {str(e)}")

            if not embeddings_list:
                mark_job_failed(job.job_id, "No valid embeddings computed")
                return

            # Stack and cache
            embeddings = np.stack(embeddings_list)
            metadata = {
                "sample_ids": sample_ids,
                "glosses": glosses,
                "signer_ids": signer_ids,
                "video_paths": video_paths,
            }
            _cache_embeddings(body.dataset_name, embeddings, metadata)

            mark_job_completed(
                job.job_id,
                result={
                    "num_embeddings": len(embeddings),
                    "embedding_dim": embeddings.shape[1],
                    "errors_count": len(errors),
                },
            )

        except JobCancelledException:
            mark_job_cancelled(job.job_id)
        except Exception as e:
            logger.exception(f"Embedding job {job.job_id} failed")
            mark_job_failed(job.job_id, str(e))

    background_tasks.add_task(run_embedding_job)

    return EmbeddingJobResponse(
        job_id=job.job_id,
        status="started",
        message="Embedding computation started",
        cached=False,
    )
