"""
Tests for the unified jobs API.

Tests cover:
- Job creation and lifecycle
- GPU monitoring
- Job cancellation
- Thread safety
"""

from __future__ import annotations

import threading
import time
from unittest.mock import Mock, patch

import pytest

# Mark all tests in this module as API tests
pytestmark = pytest.mark.api

from fastapi.testclient import TestClient

from sltk.api.main import app
from sltk.api.routers.jobs import (
    Job,
    JobCancelledException,
    JobRegistry,
    JobStatus,
    JobType,
    check_cancellation,
    check_gpu_availability,
    create_gpu_job,
    get_gpu_info,
    get_job_registry,
    is_job_cancelled,
    mark_job_cancelled,
    mark_job_completed,
    mark_job_failed,
    mark_job_started,
    release_gpu_memory,
    update_job_progress,
)

# ============================================================================
# Fixtures
# ============================================================================


@pytest.fixture
def client():
    """Create a test client."""
    return TestClient(app)


@pytest.fixture
def registry():
    """Create a fresh job registry for testing."""
    return JobRegistry()


@pytest.fixture
def reset_global_registry():
    """Reset the global registry after each test."""
    yield
    # Clear all jobs from global registry
    global_registry = get_job_registry()
    with global_registry._lock:
        global_registry._jobs.clear()
    with global_registry._gpu_lock:
        global_registry._current_gpu_job = None


# ============================================================================
# Test JobRegistry
# ============================================================================


class TestJobRegistry:
    """Test JobRegistry class."""

    def test_create_job(self, registry):
        """Test creating a new job."""
        job = registry.create_job(JobType.EMBEDDING, {"dataset": "test"})

        assert job.job_id is not None
        assert len(job.job_id) == 8
        assert job.job_type == JobType.EMBEDDING
        assert job.status == JobStatus.PENDING
        assert job.metadata["dataset"] == "test"

    def test_get_job(self, registry):
        """Test getting a job by ID."""
        job = registry.create_job(JobType.EXTRACTION)
        retrieved = registry.get_job(job.job_id)

        assert retrieved is not None
        assert retrieved.job_id == job.job_id

    def test_get_job_not_found(self, registry):
        """Test getting a non-existent job."""
        retrieved = registry.get_job("nonexistent")
        assert retrieved is None

    def test_list_jobs(self, registry):
        """Test listing jobs."""
        registry.create_job(JobType.EMBEDDING)
        registry.create_job(JobType.EXTRACTION)
        registry.create_job(JobType.SEGMENTATION)

        jobs = registry.list_jobs()
        assert len(jobs) == 3

    def test_list_jobs_filter_by_type(self, registry):
        """Test filtering jobs by type."""
        registry.create_job(JobType.EMBEDDING)
        registry.create_job(JobType.EXTRACTION)
        registry.create_job(JobType.EMBEDDING)

        embedding_jobs = registry.list_jobs(job_type=JobType.EMBEDDING)
        assert len(embedding_jobs) == 2

    def test_list_jobs_filter_by_status(self, registry):
        """Test filtering jobs by status."""
        job1 = registry.create_job(JobType.EMBEDDING)
        registry.create_job(JobType.EXTRACTION)

        registry.update_job(job1.job_id, status=JobStatus.RUNNING)

        running_jobs = registry.list_jobs(status=JobStatus.RUNNING)
        assert len(running_jobs) == 1
        assert running_jobs[0].job_id == job1.job_id

    def test_update_job(self, registry):
        """Test updating job fields."""
        job = registry.create_job(JobType.EMBEDDING)

        registry.update_job(
            job.job_id,
            status=JobStatus.RUNNING,
            progress_percent=50.0,
            current_step="Processing...",
        )

        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.RUNNING
        assert updated.progress_percent == 50.0
        assert updated.current_step == "Processing..."
        assert updated.started_at is not None

    def test_update_job_sets_completed_at(self, registry):
        """Test that completing a job sets completed_at."""
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)
        registry.update_job(job.job_id, status=JobStatus.COMPLETED)

        updated = registry.get_job(job.job_id)
        assert updated.completed_at is not None

    def test_request_cancellation(self, registry):
        """Test requesting job cancellation."""
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        result = registry.request_cancellation(job.job_id)

        assert result is True
        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.CANCELLING
        assert updated._cancel_requested is True

    def test_request_cancellation_pending_job(self, registry):
        """Test cancelling a pending job."""
        job = registry.create_job(JobType.EMBEDDING)

        result = registry.request_cancellation(job.job_id)

        assert result is True
        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.CANCELLING

    def test_request_cancellation_completed_job(self, registry):
        """Test cannot cancel a completed job."""
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.COMPLETED)

        result = registry.request_cancellation(job.job_id)

        assert result is False
        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.COMPLETED

    def test_is_cancelled(self, registry):
        """Test checking if job is cancelled."""
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        assert registry.is_cancelled(job.job_id) is False

        registry.request_cancellation(job.job_id)

        assert registry.is_cancelled(job.job_id) is True

    def test_acquire_gpu(self, registry):
        """Test acquiring GPU lock."""
        job = registry.create_job(JobType.EMBEDDING)

        result = registry.acquire_gpu(job.job_id)
        assert result is True
        assert registry.get_current_gpu_job() == job.job_id

    def test_acquire_gpu_blocked(self, registry):
        """Test GPU lock blocks second job."""
        job1 = registry.create_job(JobType.EMBEDDING)
        job2 = registry.create_job(JobType.EXTRACTION)

        registry.acquire_gpu(job1.job_id)
        registry.update_job(job1.job_id, status=JobStatus.RUNNING)

        result = registry.acquire_gpu(job2.job_id)
        assert result is False

    def test_has_active_gpu_job(self, registry):
        """Test checking for active GPU job."""
        assert registry.has_active_gpu_job() is False

        job = registry.create_job(JobType.EMBEDDING)
        registry.acquire_gpu(job.job_id)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        assert registry.has_active_gpu_job() is True

    def test_gpu_released_on_completion(self, registry):
        """Test GPU is released when job completes."""
        job = registry.create_job(JobType.EMBEDDING)
        registry.acquire_gpu(job.job_id)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        registry.update_job(job.job_id, status=JobStatus.COMPLETED)

        assert registry.get_current_gpu_job() is None

    def test_add_warning(self, registry):
        """Test adding warnings to a job."""
        job = registry.create_job(JobType.EMBEDDING)

        registry.add_warning(job.job_id, "Test warning 1")
        registry.add_warning(job.job_id, "Test warning 2")

        updated = registry.get_job(job.job_id)
        assert len(updated.warnings) == 2
        assert "Test warning 1" in updated.warnings


class TestJobRegistryThreadSafety:
    """Test thread safety of JobRegistry."""

    def test_concurrent_job_creation(self, registry):
        """Test creating jobs from multiple threads."""
        job_ids = []
        lock = threading.Lock()

        def create_job():
            job = registry.create_job(JobType.EMBEDDING)
            with lock:
                job_ids.append(job.job_id)

        threads = [threading.Thread(target=create_job) for _ in range(10)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(job_ids) == 10
        assert len(set(job_ids)) == 10  # All unique

    def test_concurrent_updates(self, registry):
        """Test updating a job from multiple threads."""
        job = registry.create_job(JobType.EMBEDDING)
        counters = [0]
        lock = threading.Lock()

        def update_progress():
            for _ in range(100):
                with lock:
                    counters[0] += 1
                    current = counters[0]
                registry.update_job(job.job_id, progress_percent=float(current))

        threads = [threading.Thread(target=update_progress) for _ in range(5)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        updated = registry.get_job(job.job_id)
        assert updated.progress_percent == 500.0  # 5 threads * 100 updates


# ============================================================================
# Test Helper Functions
# ============================================================================


class TestHelperFunctions:
    """Test helper functions for job management."""

    def test_check_gpu_availability_no_active_job(self, reset_global_registry):
        """Test GPU is available when no active job."""
        available, error = check_gpu_availability(JobType.EMBEDDING)
        # This depends on actual GPU presence
        # We just verify the function returns correctly
        assert isinstance(available, bool)
        assert error is None or isinstance(error, str)

    def test_mark_job_started(self, reset_global_registry):
        """Test marking a job as started."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)

        mark_job_started(job.job_id, total_items=100)

        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.RUNNING
        assert updated.total_items == 100
        assert updated.started_at is not None

    def test_update_job_progress(self, reset_global_registry):
        """Test updating job progress."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING, total_items=100)

        update_job_progress(job.job_id, processed_items=50, current_step="Processing")

        updated = registry.get_job(job.job_id)
        assert updated.processed_items == 50
        assert updated.current_step == "Processing"
        assert updated.progress_percent == 50.0  # Auto-calculated

    def test_mark_job_completed(self, reset_global_registry):
        """Test marking a job as completed."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        mark_job_completed(job.job_id, result={"num_embeddings": 100})

        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.COMPLETED
        assert updated.progress_percent == 100.0
        assert updated.result["num_embeddings"] == 100

    def test_mark_job_failed(self, reset_global_registry):
        """Test marking a job as failed."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        mark_job_failed(job.job_id, "Test error")

        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.FAILED
        assert updated.error == "Test error"

    def test_mark_job_cancelled(self, reset_global_registry):
        """Test marking a job as cancelled."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        mark_job_cancelled(job.job_id)

        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.CANCELLED

    def test_is_job_cancelled(self, reset_global_registry):
        """Test checking if job is cancelled."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)

        assert is_job_cancelled(job.job_id) is False

        registry.request_cancellation(job.job_id)

        assert is_job_cancelled(job.job_id) is True

    def test_check_cancellation_raises(self, reset_global_registry):
        """Test check_cancellation raises exception when cancelled."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.request_cancellation(job.job_id)

        with pytest.raises(JobCancelledException):
            check_cancellation(job.job_id)


# ============================================================================
# Test GPU Info
# ============================================================================


class TestGPUInfo:
    """Test GPU information functions."""

    def test_get_gpu_info_returns_list(self):
        """Test get_gpu_info returns a list."""
        gpus = get_gpu_info()
        assert isinstance(gpus, list)

    def test_get_gpu_info_with_gpu(self):
        """Test get_gpu_info when GPU is available."""
        # This test uses the real GPU info function
        # We just verify the structure is correct
        gpus = get_gpu_info()

        if len(gpus) > 0:
            gpu = gpus[0]
            assert hasattr(gpu, "device_id")
            assert hasattr(gpu, "name")
            assert hasattr(gpu, "total_memory_mb")
            assert hasattr(gpu, "free_memory_mb")
            assert hasattr(gpu, "used_memory_mb")
            assert gpu.total_memory_mb > 0
            assert gpu.free_memory_mb >= 0
            assert gpu.used_memory_mb >= 0

    def test_get_gpu_info_structure(self):
        """Test get_gpu_info returns correct structure."""
        gpus = get_gpu_info()

        # Should return a list (possibly empty if no GPU)
        assert isinstance(gpus, list)

        # If we have GPUs, check structure
        for gpu in gpus:
            assert hasattr(gpu, "timestamp")
            assert hasattr(gpu, "utilization_percent")

    def test_release_gpu_memory(self):
        """Test release_gpu_memory doesn't crash."""
        # Should not raise even without GPU
        release_gpu_memory()


# ============================================================================
# Test API Endpoints
# ============================================================================


class TestJobsAPIEndpoints:
    """Test jobs API endpoints."""

    def test_get_system_status(self, client, reset_global_registry):
        """Test GET /api/jobs/status endpoint."""
        response = client.get("/api/jobs/status")

        assert response.status_code == 200
        data = response.json()

        assert "gpu_available" in data
        assert "gpus" in data
        assert "has_active_gpu_job" in data
        assert "total_jobs" in data

    def test_get_gpu_status(self, client, reset_global_registry):
        """Test GET /api/jobs/gpu endpoint."""
        response = client.get("/api/jobs/gpu")

        assert response.status_code == 200
        data = response.json()

        assert isinstance(data, list)

    def test_list_jobs_empty(self, client, reset_global_registry):
        """Test GET /api/jobs/list with no jobs."""
        response = client.get("/api/jobs/list")

        assert response.status_code == 200
        data = response.json()
        assert isinstance(data, list)

    def test_list_jobs_with_filter(self, client, reset_global_registry):
        """Test GET /api/jobs/list with filters."""
        # Create a job first
        registry = get_job_registry()
        registry.create_job(JobType.EMBEDDING)

        response = client.get("/api/jobs/list?job_type=embedding")

        assert response.status_code == 200
        data = response.json()
        assert len(data) >= 1

    def test_list_jobs_invalid_filter(self, client, reset_global_registry):
        """Test GET /api/jobs/list with invalid filter."""
        response = client.get("/api/jobs/list?job_type=invalid")

        assert response.status_code == 400

    def test_get_job_not_found(self, client, reset_global_registry):
        """Test GET /api/jobs/{job_id} with non-existent job."""
        response = client.get("/api/jobs/nonexistent")

        assert response.status_code == 404

    def test_get_job_success(self, client, reset_global_registry):
        """Test GET /api/jobs/{job_id} success."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)

        response = client.get(f"/api/jobs/{job.job_id}")

        assert response.status_code == 200
        data = response.json()
        assert data["job_id"] == job.job_id
        assert data["job_type"] == "embedding"

    def test_cancel_job_not_found(self, client, reset_global_registry):
        """Test POST /api/jobs/{job_id}/cancel with non-existent job."""
        response = client.post("/api/jobs/nonexistent/cancel")

        assert response.status_code == 404

    def test_cancel_job_success(self, client, reset_global_registry):
        """Test POST /api/jobs/{job_id}/cancel success."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.RUNNING)

        response = client.post(f"/api/jobs/{job.job_id}/cancel")

        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "cancelling"

        # Verify job was marked for cancellation
        updated = registry.get_job(job.job_id)
        assert updated.status == JobStatus.CANCELLING

    def test_cancel_completed_job(self, client, reset_global_registry):
        """Test cancelling a completed job."""
        registry = get_job_registry()
        job = registry.create_job(JobType.EMBEDDING)
        registry.update_job(job.job_id, status=JobStatus.COMPLETED)

        response = client.post(f"/api/jobs/{job.job_id}/cancel")

        assert response.status_code == 200
        data = response.json()
        assert "cannot be cancelled" in data["message"].lower()

    def test_force_release_gpu(self, client, reset_global_registry):
        """Test POST /api/jobs/gpu/release endpoint."""
        response = client.post("/api/jobs/gpu/release")

        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "released"


# ============================================================================
# Test Job Lifecycle
# ============================================================================


class TestJobLifecycle:
    """Test complete job lifecycle scenarios."""

    def test_embedding_job_lifecycle(self, reset_global_registry):
        """Test complete embedding job lifecycle."""
        registry = get_job_registry()

        # Create job
        job = registry.create_job(
            JobType.EMBEDDING,
            metadata={"dataset": "test_dataset"},
        )
        assert job.status == JobStatus.PENDING

        # Acquire GPU
        assert registry.acquire_gpu(job.job_id) is True

        # Start job
        mark_job_started(job.job_id, total_items=100)
        job = registry.get_job(job.job_id)
        assert job.status == JobStatus.RUNNING
        assert job.started_at is not None

        # Update progress
        for i in range(1, 101):
            update_job_progress(
                job.job_id,
                processed_items=i,
                current_step=f"Processing item {i}",
            )

        job = registry.get_job(job.job_id)
        assert job.progress_percent == 100.0

        # Complete job
        mark_job_completed(job.job_id, result={"num_embeddings": 100})

        job = registry.get_job(job.job_id)
        assert job.status == JobStatus.COMPLETED
        assert job.completed_at is not None
        assert job.result["num_embeddings"] == 100

        # GPU should be released
        assert registry.get_current_gpu_job() is None

    def test_job_cancellation_lifecycle(self, reset_global_registry):
        """Test job cancellation lifecycle."""
        registry = get_job_registry()

        # Create and start job
        job = registry.create_job(JobType.EXTRACTION)
        registry.acquire_gpu(job.job_id)
        mark_job_started(job.job_id, total_items=100)

        # Request cancellation
        registry.request_cancellation(job.job_id)

        job = registry.get_job(job.job_id)
        assert job.status == JobStatus.CANCELLING
        assert job._cancel_requested is True

        # Job checks cancellation and marks itself cancelled
        if is_job_cancelled(job.job_id):
            mark_job_cancelled(job.job_id)

        job = registry.get_job(job.job_id)
        assert job.status == JobStatus.CANCELLED

        # GPU should be released
        assert registry.get_current_gpu_job() is None

    def test_job_failure_lifecycle(self, reset_global_registry):
        """Test job failure lifecycle."""
        registry = get_job_registry()

        # Create and start job
        job = registry.create_job(JobType.SEGMENTATION)
        registry.acquire_gpu(job.job_id)
        mark_job_started(job.job_id, total_items=50)

        # Simulate failure
        mark_job_failed(job.job_id, "Model not found")

        job = registry.get_job(job.job_id)
        assert job.status == JobStatus.FAILED
        assert job.error == "Model not found"
        assert job.completed_at is not None

        # GPU should be released
        assert registry.get_current_gpu_job() is None
