"""Tests for M2: Job executor, GPU allocator, and jobs API endpoints."""

from __future__ import annotations

import json
from unittest.mock import patch, AsyncMock, MagicMock

from radix_edge.executor.gpu_allocator import (
    get_available_gpus,
    allocate,
    release,
    _find_nvlink_group,
)


# --- GPU Allocator ---

class TestGpuAllocator:
    def _insert_gpus(self, db, n=4, nvlink_peers=None):
        """Insert n idle GPUs. nvlink_peers: dict mapping gpu_index -> list of peers."""
        for i in range(n):
            peers = json.dumps(nvlink_peers.get(i, [])) if nvlink_peers else "[]"
            db.execute(
                """INSERT INTO gpus (gpu_index, name, memory_total_mb, nvlink_peers, status)
                VALUES (?, ?, ?, ?, ?)""",
                (i, "A100", 81920, peers, "idle"),
            )
        db.commit()

    def test_get_available_gpus(self, db):
        self._insert_gpus(db, n=4)
        # Mark GPU 1 as allocated
        db.execute("UPDATE gpus SET status = 'allocated', assigned_job_id = 'job-x' WHERE gpu_index = 1")
        db.commit()

        available = get_available_gpus(db)
        assert available == [0, 2, 3]

    def test_allocate_single_gpu(self, db):
        self._insert_gpus(db, n=4)
        indices = allocate("job-1", 1)
        assert indices is not None
        assert len(indices) == 1

        # Verify GPU is marked allocated
        row = db.execute("SELECT * FROM gpus WHERE gpu_index = ?", (indices[0],)).fetchone()
        assert row["status"] == "allocated"
        assert row["assigned_job_id"] == "job-1"

    def test_allocate_multi_gpu(self, db):
        self._insert_gpus(db, n=4)
        indices = allocate("job-1", 2)
        assert indices is not None
        assert len(indices) == 2

    def test_allocate_insufficient_gpus(self, db):
        self._insert_gpus(db, n=2)
        # Allocate all 2
        allocate("job-1", 2)
        # Try to allocate 1 more — should fail
        result = allocate("job-2", 1)
        assert result is None

    def test_allocate_topology_aware(self, db):
        """Multi-GPU allocation prefers NVLink-connected GPUs."""
        nvlink = {0: [1], 1: [0], 2: [3], 3: [2]}
        self._insert_gpus(db, n=4, nvlink_peers=nvlink)

        indices = allocate("job-1", 2)
        assert indices is not None
        assert len(indices) == 2
        # Should pick a NVLink-connected pair
        assert (indices == [0, 1]) or (indices == [2, 3])

    def test_release(self, db):
        self._insert_gpus(db, n=2)
        allocate("job-1", 2)

        # Verify both allocated
        available = get_available_gpus(db)
        assert available == []

        # Release
        freed = release("job-1")
        assert sorted(freed) == [0, 1]

        # Verify back to idle
        available = get_available_gpus(db)
        assert sorted(available) == [0, 1]

    def test_release_nonexistent_job(self, db):
        self._insert_gpus(db, n=2)
        freed = release("job-nonexistent")
        assert freed == []


class TestNvlinkGroupFinding:
    def test_find_nvlink_pair(self):
        gpus = [
            {"gpu_index": 0, "nvlink_peers": "[1]"},
            {"gpu_index": 1, "nvlink_peers": "[0]"},
            {"gpu_index": 2, "nvlink_peers": "[]"},
            {"gpu_index": 3, "nvlink_peers": "[]"},
        ]
        result = _find_nvlink_group(gpus, 2, {0, 1, 2, 3})
        assert result == [0, 1]

    def test_no_nvlink_returns_none(self):
        gpus = [
            {"gpu_index": 0, "nvlink_peers": "[]"},
            {"gpu_index": 1, "nvlink_peers": "[]"},
        ]
        result = _find_nvlink_group(gpus, 2, {0, 1})
        assert result is None

    def test_single_gpu_returns_none(self):
        gpus = [{"gpu_index": 0, "nvlink_peers": "[1]"}]
        result = _find_nvlink_group(gpus, 1, {0})
        assert result is None  # No topology preference for single GPU


# --- Docker Command Building ---

class TestDockerCommand:
    def test_basic_command(self):
        from radix_edge.executor.executor import _build_docker_command

        cmd = _build_docker_command(
            image="nvidia/cuda:12.4",
            command=["python", "train.py"],
            args=["--epochs", "10"],
            env={"BATCH_SIZE": "32"},
            gpu_indices=[0],
            num_gpus=1,
            job_dir="/tmp/job-123",
            job_id="job-test123",
            timeout_seconds=3600,
        )

        assert "docker" == cmd[0]
        assert "run" == cmd[1]
        assert "--rm" in cmd
        assert "-v" in cmd
        assert "/tmp/job-123:/workspace" in cmd
        assert "nvidia/cuda:12.4" in cmd
        assert "python" in cmd
        assert "train.py" in cmd
        assert "--epochs" in cmd
        assert "10" in cmd
        assert "CUDA_VISIBLE_DEVICES=0" in " ".join(cmd)
        assert "BATCH_SIZE=32" in " ".join(cmd)

    def test_multi_gpu_wraps_torchrun(self):
        from radix_edge.executor.executor import _build_docker_command

        cmd = _build_docker_command(
            image="myimg:latest",
            command=["python", "train.py"],
            args=None,
            env=None,
            gpu_indices=[0, 1, 2, 3],
            num_gpus=4,
            job_dir="/tmp/job-456",
            job_id="job-test456",
            timeout_seconds=3600,
        )

        assert "torchrun" in cmd
        assert "--nproc_per_node=4" in cmd
        assert "--standalone" in cmd
        assert "CUDA_VISIBLE_DEVICES=0,1,2,3" in " ".join(cmd)

    def test_no_gpu(self):
        from radix_edge.executor.executor import _build_docker_command

        cmd = _build_docker_command(
            image="myimg:latest",
            command=["echo", "hello"],
            args=None,
            env=None,
            gpu_indices=[],
            num_gpus=0,
            job_dir="/tmp/job-789",
            job_id="job-test789",
            timeout_seconds=60,
        )

        assert "--gpus" not in cmd
        assert "CUDA_VISIBLE_DEVICES" not in " ".join(cmd)


# --- Jobs API Endpoints ---

class TestJobsApi:
    def test_submit_job(self, client, db):
        resp = client.post("/v1/jobs", json={
            "image": "nvidia/cuda:12.4",
            "user_id": "test-user",
        })
        assert resp.status_code == 201
        data = resp.json()
        assert data["status"] == "queued"
        assert data["job_id"].startswith("job-")

        # Verify in database
        row = db.execute("SELECT * FROM jobs WHERE job_id = ?", (data["job_id"],)).fetchone()
        assert row["image"] == "nvidia/cuda:12.4"
        assert row["user_id"] == "test-user"
        assert row["status"] == "queued"
        assert row["priority"] == 5

    def test_submit_job_with_all_fields(self, client, db):
        resp = client.post("/v1/jobs", json={
            "image": "myimg:latest",
            "command": ["python", "train.py"],
            "args": ["--epochs", "10"],
            "env": {"BATCH_SIZE": "32"},
            "name": "Training Run",
            "job_type": "training",
            "priority": 8,
            "num_gpus": 2,
            "user_id": "user-a",
            "timeout_seconds": 7200,
        })
        assert resp.status_code == 201
        data = resp.json()

        row = db.execute("SELECT * FROM jobs WHERE job_id = ?", (data["job_id"],)).fetchone()
        assert json.loads(row["command"]) == ["python", "train.py"]
        assert row["priority"] == 8
        assert row["num_gpus"] == 2
        assert row["job_type"] == "training"

    def test_list_jobs_empty(self, client, db):
        resp = client.get("/v1/jobs")
        assert resp.status_code == 200
        assert resp.json()["jobs"] == []

    def test_list_jobs_with_filter(self, client, db):
        # Submit two jobs
        client.post("/v1/jobs", json={"image": "img:1", "user_id": "user-a"})
        client.post("/v1/jobs", json={"image": "img:2", "user_id": "user-b"})

        # Filter by user
        resp = client.get("/v1/jobs?user_id=user-a")
        assert resp.status_code == 200
        jobs = resp.json()["jobs"]
        assert len(jobs) == 1
        assert jobs[0]["user_id"] == "user-a"

    def test_get_job(self, client, db):
        resp = client.post("/v1/jobs", json={"image": "img:1", "user_id": "u1"})
        job_id = resp.json()["job_id"]

        resp = client.get(f"/v1/jobs/{job_id}")
        assert resp.status_code == 200
        assert resp.json()["job_id"] == job_id
        assert resp.json()["status"] == "queued"

    def test_get_job_not_found(self, client, db):
        resp = client.get("/v1/jobs/nonexistent")
        assert resp.status_code == 404

    def test_cancel_queued_job(self, client, db):
        resp = client.post("/v1/jobs", json={"image": "img:1", "user_id": "u1"})
        job_id = resp.json()["job_id"]

        resp = client.delete(f"/v1/jobs/{job_id}")
        assert resp.status_code == 200
        assert resp.json()["status"] == "cancelled"

        # Verify in DB
        row = db.execute("SELECT status FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
        assert row["status"] == "cancelled"

    def test_cancel_already_terminal(self, client, db):
        resp = client.post("/v1/jobs", json={"image": "img:1", "user_id": "u1"})
        job_id = resp.json()["job_id"]

        # Cancel once
        client.delete(f"/v1/jobs/{job_id}")
        # Cancel again — should fail
        resp = client.delete(f"/v1/jobs/{job_id}")
        assert resp.status_code == 409

    def test_cancel_not_found(self, client, db):
        resp = client.delete("/v1/jobs/nonexistent")
        assert resp.status_code == 404

    def test_submit_validation(self, client, db):
        # Missing required field
        resp = client.post("/v1/jobs", json={"user_id": "u1"})
        assert resp.status_code == 422

        # Invalid priority
        resp = client.post("/v1/jobs", json={"image": "img", "priority": 99})
        assert resp.status_code == 422
