"""
Unified job management API for SLTK.

Provides centralized job tracking for all GPU-intensive tasks:
- Pose extraction (MediaPipe, WiLoR, NLF)
- Embedding computation (SignRep)
- Segmentation

Features:
- GPU monitoring and usage tracking
- Single GPU task enforcement
- Job cancellation with graceful cleanup
- Progress streaming via WebSocket or polling
"""

from __future__ import annotations

import asyncio
import gc
import logging
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any

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

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


# ============================================================================
# Job Types and Status
# ============================================================================


class JobType(str, Enum):
    """Types of GPU-intensive jobs."""

    EXTRACTION = "extraction"
    EMBEDDING = "embedding"
    SEGMENTATION = "segmentation"
    SIGNREP_DICTIONARY = "signrep_dictionary"
    SIGNREP_CONTINUOUS = "signrep_continuous"
    SIGNREP_SPOTTING = "signrep_spotting"


class JobStatus(str, Enum):
    """Job status states."""

    PENDING = "pending"
    QUEUED = "queued"
    RUNNING = "running"
    CANCELLING = "cancelling"
    CANCELLED = "cancelled"
    COMPLETED = "completed"
    FAILED = "failed"


# ============================================================================
# GPU Monitoring
# ============================================================================


@dataclass
class GPUInfo:
    """GPU information snapshot."""

    device_id: int = 0
    name: str = "Unknown"
    total_memory_mb: float = 0
    used_memory_mb: float = 0
    free_memory_mb: float = 0
    utilization_percent: float = 0
    temperature_celsius: float | None = None
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


def get_gpu_info() -> list[GPUInfo]:
    """
    Get current GPU information.

    Returns list of GPUInfo for all available GPUs.
    Returns empty list if no GPU or pynvml unavailable.
    """
    try:
        import torch

        if not torch.cuda.is_available():
            return []

        gpus = []
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            mem_info = torch.cuda.mem_get_info(i)

            gpu = GPUInfo(
                device_id=i,
                name=props.name,
                total_memory_mb=props.total_memory / (1024 * 1024),
                free_memory_mb=mem_info[0] / (1024 * 1024),
                used_memory_mb=(props.total_memory - mem_info[0]) / (1024 * 1024),
            )

            # Try to get utilization via pynvml
            try:
                import pynvml

                pynvml.nvmlInit()
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                gpu.utilization_percent = util.gpu
                temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
                gpu.temperature_celsius = temp
                pynvml.nvmlShutdown()
            except (ImportError, Exception):
                pass

            gpus.append(gpu)

        return gpus

    except Exception as e:
        logger.debug(f"Failed to get GPU info: {e}")
        return []


def release_gpu_memory():
    """Force release of GPU memory."""
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
            gc.collect()
            logger.info("GPU memory released")
    except ImportError:
        pass
    except Exception as e:
        logger.warning(f"Failed to release GPU memory: {e}")


# ============================================================================
# Job Registry (Thread-Safe)
# ============================================================================


@dataclass
class Job:
    """A tracked job with progress and status."""

    job_id: str
    job_type: JobType
    status: JobStatus = JobStatus.PENDING
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    started_at: str | None = None
    completed_at: str | None = None

    # Progress tracking
    progress_percent: float = 0.0
    current_step: str = ""
    total_items: int = 0
    processed_items: int = 0

    # Results
    result: dict[str, Any] = field(default_factory=dict)
    error: str | None = None
    warnings: list[str] = field(default_factory=list)

    # Metadata
    metadata: dict[str, Any] = field(default_factory=dict)

    # Internal: cancellation flag
    _cancel_requested: bool = field(default=False, repr=False)


class JobRegistry:
    """Thread-safe job registry with GPU exclusivity enforcement."""

    def __init__(self, max_jobs: int = 100, ttl_hours: int = 24):
        self._jobs: dict[str, Job] = {}
        self._lock = threading.Lock()
        self._gpu_lock = threading.Lock()
        self._current_gpu_job: str | None = None
        self._max_jobs = max_jobs
        self._ttl_hours = ttl_hours

        # WebSocket connections for progress updates
        self._websocket_connections: dict[str, list[WebSocket]] = {}
        self._ws_lock = threading.Lock()

    def create_job(self, job_type: JobType, metadata: dict[str, Any] | None = None) -> Job:
        """Create a new job and register it."""
        job_id = str(uuid.uuid4())[:8]

        job = Job(
            job_id=job_id,
            job_type=job_type,
            status=JobStatus.PENDING,
            metadata=metadata or {},
        )

        with self._lock:
            self._cleanup_old_jobs()
            self._jobs[job_id] = job

        return job

    def get_job(self, job_id: str) -> Job | None:
        """Get a job by ID."""
        with self._lock:
            return self._jobs.get(job_id)

    def list_jobs(
        self,
        job_type: JobType | None = None,
        status: JobStatus | None = None,
        limit: int = 50,
    ) -> list[Job]:
        """List jobs with optional filtering."""
        with self._lock:
            jobs = list(self._jobs.values())

        if job_type:
            jobs = [j for j in jobs if j.job_type == job_type]
        if status:
            jobs = [j for j in jobs if j.status == status]

        # Sort by created_at descending
        jobs.sort(key=lambda j: j.created_at, reverse=True)
        return jobs[:limit]

    def update_job(
        self,
        job_id: str,
        status: JobStatus | None = None,
        progress_percent: float | None = None,
        current_step: str | None = None,
        processed_items: int | None = None,
        total_items: int | None = None,
        error: str | None = None,
        result: dict[str, Any] | None = None,
    ) -> bool:
        """Update job fields. Returns True if job exists."""
        with self._lock:
            if job_id not in self._jobs:
                return False

            job = self._jobs[job_id]

            if status is not None:
                old_status = job.status
                job.status = status

                if status == JobStatus.RUNNING and job.started_at is None:
                    job.started_at = datetime.now().isoformat()
                elif status in (
                    JobStatus.COMPLETED,
                    JobStatus.FAILED,
                    JobStatus.CANCELLED,
                ):
                    job.completed_at = datetime.now().isoformat()

                # Release GPU lock when job completes or is cancelled
                if old_status in (JobStatus.RUNNING, JobStatus.CANCELLING) and status in (
                    JobStatus.COMPLETED,
                    JobStatus.FAILED,
                    JobStatus.CANCELLED,
                ):
                    self._release_gpu_lock(job_id)

            if progress_percent is not None:
                job.progress_percent = progress_percent
            if current_step is not None:
                job.current_step = current_step
            if processed_items is not None:
                job.processed_items = processed_items
            if total_items is not None:
                job.total_items = total_items
            if error is not None:
                job.error = error
            if result is not None:
                job.result = result

            return True

    def add_warning(self, job_id: str, warning: str) -> bool:
        """Add a warning message to a job."""
        with self._lock:
            if job_id not in self._jobs:
                return False
            self._jobs[job_id].warnings.append(warning)
            return True

    def request_cancellation(self, job_id: str) -> bool:
        """
        Request job cancellation.

        Sets the cancel flag and status to CANCELLING.
        The actual job implementation must check is_cancelled() periodically.
        """
        with self._lock:
            if job_id not in self._jobs:
                return False

            job = self._jobs[job_id]

            # Can only cancel pending/queued/running jobs
            if job.status not in (
                JobStatus.PENDING,
                JobStatus.QUEUED,
                JobStatus.RUNNING,
            ):
                return False

            job._cancel_requested = True
            job.status = JobStatus.CANCELLING

            return True

    def is_cancelled(self, job_id: str) -> bool:
        """Check if cancellation has been requested for a job."""
        with self._lock:
            if job_id not in self._jobs:
                return False
            return self._jobs[job_id]._cancel_requested

    def acquire_gpu(self, job_id: str) -> bool:
        """
        Try to acquire GPU lock for a job.

        Returns True if acquired, False if another GPU job is running.
        """
        with self._gpu_lock:
            if self._current_gpu_job is not None:
                # Check if the current GPU job is still active
                current_job = self._jobs.get(self._current_gpu_job)
                if current_job and current_job.status == JobStatus.RUNNING:
                    return False
                # Previous job finished, clear the lock
                self._current_gpu_job = None

            self._current_gpu_job = job_id
            return True

    def _release_gpu_lock(self, job_id: str):
        """Release GPU lock for a job."""
        with self._gpu_lock:
            if self._current_gpu_job == job_id:
                self._current_gpu_job = None
                release_gpu_memory()

    def get_current_gpu_job(self) -> str | None:
        """Get the ID of the currently running GPU job."""
        with self._gpu_lock:
            return self._current_gpu_job

    def has_active_gpu_job(self) -> bool:
        """Check if a GPU job is currently running."""
        with self._gpu_lock:
            if self._current_gpu_job is None:
                return False
            job = self._jobs.get(self._current_gpu_job)
            return job is not None and job.status == JobStatus.RUNNING

    def _cleanup_old_jobs(self):
        """Remove old completed jobs (called within lock)."""
        from datetime import timedelta

        cutoff = datetime.now() - timedelta(hours=self._ttl_hours)
        cutoff_str = cutoff.isoformat()

        to_remove = []
        for job_id, job in self._jobs.items():
            if job.status in (
                JobStatus.COMPLETED,
                JobStatus.FAILED,
                JobStatus.CANCELLED,
            ):
                if job.completed_at and job.completed_at < cutoff_str:
                    to_remove.append(job_id)

        for job_id in to_remove:
            del self._jobs[job_id]

        # Also enforce max jobs limit
        if len(self._jobs) > self._max_jobs:
            completed_jobs = [
                (jid, j)
                for jid, j in self._jobs.items()
                if j.status in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED)
            ]
            completed_jobs.sort(key=lambda x: x[1].completed_at or "")

            excess = len(self._jobs) - self._max_jobs
            for job_id, _ in completed_jobs[:excess]:
                del self._jobs[job_id]

    # WebSocket support for real-time progress
    async def register_websocket(self, job_id: str, websocket: WebSocket):
        """Register a WebSocket connection for job progress updates."""
        with self._ws_lock:
            if job_id not in self._websocket_connections:
                self._websocket_connections[job_id] = []
            self._websocket_connections[job_id].append(websocket)

    async def unregister_websocket(self, job_id: str, websocket: WebSocket):
        """Unregister a WebSocket connection."""
        with self._ws_lock:
            if job_id in self._websocket_connections:
                try:
                    self._websocket_connections[job_id].remove(websocket)
                except ValueError:
                    pass

    async def broadcast_progress(self, job_id: str, data: dict):
        """Broadcast progress update to all connected WebSocket clients."""
        with self._ws_lock:
            connections = self._websocket_connections.get(job_id, [])

        for ws in connections:
            try:
                await ws.send_json(data)
            except Exception:
                pass


# Global job registry
_job_registry = JobRegistry()


def get_job_registry() -> JobRegistry:
    """Get the global job registry."""
    return _job_registry


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


class GPUInfoResponse(BaseModel):
    """GPU information response."""

    device_id: int
    name: str
    total_memory_mb: float
    used_memory_mb: float
    free_memory_mb: float
    utilization_percent: float
    temperature_celsius: float | None = None
    timestamp: str


class JobResponse(BaseModel):
    """Job information response."""

    job_id: str
    job_type: str
    status: str
    created_at: str
    started_at: str | None = None
    completed_at: str | None = None
    progress_percent: float = 0.0
    current_step: str = ""
    total_items: int = 0
    processed_items: int = 0
    error: str | None = None
    warnings: list[str] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)
    result: dict[str, Any] = Field(default_factory=dict)


class SystemStatusResponse(BaseModel):
    """System status including GPU and job information."""

    gpu_available: bool
    gpus: list[GPUInfoResponse] = Field(default_factory=list)
    has_active_gpu_job: bool
    current_gpu_job_id: str | None = None
    current_gpu_job: JobResponse | None = None
    total_jobs: int = 0
    pending_jobs: int = 0
    running_jobs: int = 0


class CancelJobResponse(BaseModel):
    """Response from cancelling a job."""

    job_id: str
    status: str
    message: str


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


@router.get("/status", response_model=SystemStatusResponse)
async def get_system_status():
    """
    Get system status including GPU info and active jobs.

    Returns GPU availability, memory usage, and current job information.
    """
    registry = get_job_registry()

    # Get GPU info
    gpus = get_gpu_info()
    gpu_responses = [
        GPUInfoResponse(
            device_id=g.device_id,
            name=g.name,
            total_memory_mb=round(g.total_memory_mb, 1),
            used_memory_mb=round(g.used_memory_mb, 1),
            free_memory_mb=round(g.free_memory_mb, 1),
            utilization_percent=g.utilization_percent,
            temperature_celsius=g.temperature_celsius,
            timestamp=g.timestamp,
        )
        for g in gpus
    ]

    # Get job counts
    all_jobs = registry.list_jobs(limit=1000)
    pending_count = len([j for j in all_jobs if j.status == JobStatus.PENDING])
    running_count = len([j for j in all_jobs if j.status == JobStatus.RUNNING])

    # Get current GPU job
    current_gpu_job_id = registry.get_current_gpu_job()
    current_gpu_job = None
    if current_gpu_job_id:
        job = registry.get_job(current_gpu_job_id)
        if job:
            current_gpu_job = JobResponse(
                job_id=job.job_id,
                job_type=job.job_type.value,
                status=job.status.value,
                created_at=job.created_at,
                started_at=job.started_at,
                completed_at=job.completed_at,
                progress_percent=job.progress_percent,
                current_step=job.current_step,
                total_items=job.total_items,
                processed_items=job.processed_items,
                error=job.error,
                warnings=job.warnings,
                metadata=job.metadata,
                result=job.result,
            )

    return SystemStatusResponse(
        gpu_available=len(gpus) > 0,
        gpus=gpu_responses,
        has_active_gpu_job=registry.has_active_gpu_job(),
        current_gpu_job_id=current_gpu_job_id,
        current_gpu_job=current_gpu_job,
        total_jobs=len(all_jobs),
        pending_jobs=pending_count,
        running_jobs=running_count,
    )


@router.get("/gpu", response_model=list[GPUInfoResponse])
async def get_gpu_status():
    """
    Get detailed GPU information.

    Returns memory usage, utilization, and temperature for all GPUs.
    """
    gpus = get_gpu_info()

    return [
        GPUInfoResponse(
            device_id=g.device_id,
            name=g.name,
            total_memory_mb=round(g.total_memory_mb, 1),
            used_memory_mb=round(g.used_memory_mb, 1),
            free_memory_mb=round(g.free_memory_mb, 1),
            utilization_percent=g.utilization_percent,
            temperature_celsius=g.temperature_celsius,
            timestamp=g.timestamp,
        )
        for g in gpus
    ]


@router.get("/list", response_model=list[JobResponse])
async def list_jobs(
    job_type: str | None = None,
    status: str | None = None,
    limit: int = 50,
):
    """
    List jobs with optional filtering.

    Args:
        job_type: Filter by job type (extraction, embedding, segmentation)
        status: Filter by status (pending, running, completed, failed, cancelled)
        limit: Maximum number of jobs to return
    """
    registry = get_job_registry()

    # Convert string filters to enums
    type_filter = None
    if job_type:
        try:
            type_filter = JobType(job_type)
        except ValueError:
            raise HTTPException(400, f"Invalid job_type: {job_type}")

    status_filter = None
    if status:
        try:
            status_filter = JobStatus(status)
        except ValueError:
            raise HTTPException(400, f"Invalid status: {status}")

    jobs = registry.list_jobs(job_type=type_filter, status=status_filter, limit=limit)

    return [
        JobResponse(
            job_id=j.job_id,
            job_type=j.job_type.value,
            status=j.status.value,
            created_at=j.created_at,
            started_at=j.started_at,
            completed_at=j.completed_at,
            progress_percent=j.progress_percent,
            current_step=j.current_step,
            total_items=j.total_items,
            processed_items=j.processed_items,
            error=j.error,
            warnings=j.warnings,
            metadata=j.metadata,
            result=j.result,
        )
        for j in jobs
    ]


@router.get("/{job_id}", response_model=JobResponse)
async def get_job(job_id: str):
    """Get detailed information about a specific job."""
    registry = get_job_registry()
    job = registry.get_job(job_id)

    if not job:
        raise HTTPException(404, f"Job not found: {job_id}")

    return JobResponse(
        job_id=job.job_id,
        job_type=job.job_type.value,
        status=job.status.value,
        created_at=job.created_at,
        started_at=job.started_at,
        completed_at=job.completed_at,
        progress_percent=job.progress_percent,
        current_step=job.current_step,
        total_items=job.total_items,
        processed_items=job.processed_items,
        error=job.error,
        warnings=job.warnings,
        metadata=job.metadata,
        result=job.result,
    )


@router.post("/{job_id}/cancel", response_model=CancelJobResponse)
async def cancel_job(job_id: str):
    """
    Cancel a running or pending job.

    For GPU jobs, this will:
    1. Set the job's cancel flag
    2. Wait for the job to acknowledge cancellation
    3. Release GPU memory

    The actual job must check for cancellation periodically.
    Jobs may take a moment to cancel at a safe checkpoint.
    """
    registry = get_job_registry()
    job = registry.get_job(job_id)

    if not job:
        raise HTTPException(404, f"Job not found: {job_id}")

    # Check if job can be cancelled
    if job.status not in (
        JobStatus.PENDING,
        JobStatus.QUEUED,
        JobStatus.RUNNING,
    ):
        return CancelJobResponse(
            job_id=job_id,
            status=job.status.value,
            message=f"Job cannot be cancelled (status: {job.status.value})",
        )

    # Request cancellation
    success = registry.request_cancellation(job_id)

    if not success:
        return CancelJobResponse(
            job_id=job_id,
            status=job.status.value,
            message="Failed to request cancellation",
        )

    # Force GPU memory release if this is a GPU job
    if job.job_type in (
        JobType.EXTRACTION,
        JobType.EMBEDDING,
        JobType.SEGMENTATION,
        JobType.SIGNREP_DICTIONARY,
        JobType.SIGNREP_CONTINUOUS,
        JobType.SIGNREP_SPOTTING,
    ):
        release_gpu_memory()

    return CancelJobResponse(
        job_id=job_id,
        status="cancelling",
        message="Cancellation requested. Job will stop at next checkpoint.",
    )


@router.post("/gpu/release")
async def force_release_gpu():
    """
    Force release GPU memory.

    Clears CUDA cache and runs garbage collection.
    Use when GPU memory seems stuck after a failed job.
    """
    release_gpu_memory()
    return {"status": "released", "message": "GPU memory cache cleared"}


@router.websocket("/{job_id}/progress")
async def job_progress_websocket(websocket: WebSocket, job_id: str):
    """
    WebSocket endpoint for real-time job progress updates.

    Connect to receive progress updates as the job runs.
    Messages are JSON objects with current job state.
    """
    registry = get_job_registry()
    job = registry.get_job(job_id)

    if not job:
        await websocket.close(code=4004, reason="Job not found")
        return

    await websocket.accept()
    await registry.register_websocket(job_id, websocket)

    try:
        # Send initial state
        await websocket.send_json(
            {
                "type": "initial",
                "job_id": job.job_id,
                "status": job.status.value,
                "progress_percent": job.progress_percent,
                "current_step": job.current_step,
                "total_items": job.total_items,
                "processed_items": job.processed_items,
            }
        )

        # Keep connection alive and send updates
        while True:
            # Poll for updates every 500ms
            await asyncio.sleep(0.5)

            job = registry.get_job(job_id)
            if not job:
                break

            # Send progress update
            await websocket.send_json(
                {
                    "type": "progress",
                    "job_id": job.job_id,
                    "status": job.status.value,
                    "progress_percent": job.progress_percent,
                    "current_step": job.current_step,
                    "total_items": job.total_items,
                    "processed_items": job.processed_items,
                }
            )

            # Exit if job is complete
            if job.status in (
                JobStatus.COMPLETED,
                JobStatus.FAILED,
                JobStatus.CANCELLED,
            ):
                # Send final state
                await websocket.send_json(
                    {
                        "type": "complete",
                        "job_id": job.job_id,
                        "status": job.status.value,
                        "result": job.result,
                        "error": job.error,
                    }
                )
                break

    except WebSocketDisconnect:
        pass
    finally:
        await registry.unregister_websocket(job_id, websocket)


# ============================================================================
# Helper functions for other routers
# ============================================================================


def check_gpu_availability(job_type: JobType) -> tuple[bool, str | None]:
    """
    Check if GPU is available for a new job.

    Returns:
        Tuple of (available, error_message)
    """
    registry = get_job_registry()

    if registry.has_active_gpu_job():
        current_job_id = registry.get_current_gpu_job()
        return (
            False,
            f"GPU busy with job {current_job_id}. " "Wait for it to complete or cancel it first.",
        )

    # Check GPU memory
    gpus = get_gpu_info()
    if not gpus:
        return False, "No GPU available"

    # Require at least 2GB free
    min_free_mb = 2000
    if all(g.free_memory_mb < min_free_mb for g in gpus):
        return (
            False,
            f"Insufficient GPU memory. Need at least {min_free_mb}MB free.",
        )

    return True, None


def create_gpu_job(
    job_type: JobType,
    metadata: dict[str, Any] | None = None,
) -> tuple[Job | None, str | None]:
    """
    Create a new GPU job if GPU is available.

    Returns:
        Tuple of (job, error_message)
    """
    available, error = check_gpu_availability(job_type)
    if not available:
        return None, error

    registry = get_job_registry()
    job = registry.create_job(job_type, metadata)

    # Try to acquire GPU
    if not registry.acquire_gpu(job.job_id):
        return None, "Failed to acquire GPU lock"

    return job, None


def mark_job_started(job_id: str, total_items: int = 0):
    """Mark a job as started."""
    registry = get_job_registry()
    registry.update_job(
        job_id,
        status=JobStatus.RUNNING,
        total_items=total_items,
    )


def update_job_progress(
    job_id: str,
    processed_items: int,
    current_step: str = "",
    progress_percent: float | None = None,
):
    """Update job progress."""
    registry = get_job_registry()
    job = registry.get_job(job_id)

    if job and progress_percent is None and job.total_items > 0:
        progress_percent = (processed_items / job.total_items) * 100

    registry.update_job(
        job_id,
        processed_items=processed_items,
        current_step=current_step,
        progress_percent=progress_percent or 0,
    )


def mark_job_completed(job_id: str, result: dict[str, Any] | None = None):
    """Mark a job as completed."""
    registry = get_job_registry()
    registry.update_job(
        job_id,
        status=JobStatus.COMPLETED,
        progress_percent=100.0,
        result=result or {},
    )


def mark_job_failed(job_id: str, error: str):
    """Mark a job as failed."""
    registry = get_job_registry()
    registry.update_job(job_id, status=JobStatus.FAILED, error=error)
    release_gpu_memory()


def mark_job_cancelled(job_id: str):
    """Mark a job as cancelled."""
    registry = get_job_registry()
    registry.update_job(job_id, status=JobStatus.CANCELLED)
    release_gpu_memory()


def is_job_cancelled(job_id: str) -> bool:
    """Check if a job has been cancelled."""
    return get_job_registry().is_cancelled(job_id)


class JobCancelledException(Exception):
    """Exception raised when a job is cancelled."""

    pass


def check_cancellation(job_id: str):
    """
    Check if job is cancelled and raise if so.

    Call this periodically in long-running tasks.
    """
    if is_job_cancelled(job_id):
        release_gpu_memory()
        raise JobCancelledException(f"Job {job_id} was cancelled")
