"""Tests for M1: GPU observer, topology detection, and GPU API endpoints."""

from __future__ import annotations

import json
from unittest.mock import patch, MagicMock

from radix_edge.observer.gpu_observer import (
    poll_gpus,
    store_gpu_state,
    prune_old_observations,
    compute_observation_vector,
    _parse_mib,
    _parse_float,
)
from radix_edge.observer.gpu_topology import parse_topology_output


# --- nvidia-smi parsing helpers ---

class TestParsers:
    def test_parse_mib_normal(self):
        assert _parse_mib("81920 MiB") == 81920

    def test_parse_mib_na(self):
        assert _parse_mib("[N/A]") == 0

    def test_parse_mib_empty(self):
        assert _parse_mib("") == 0

    def test_parse_float_percent(self):
        assert _parse_float("45.2 %") == 45.2

    def test_parse_float_watts(self):
        assert _parse_float("250.00 W") == 250.0

    def test_parse_float_na(self):
        assert _parse_float("[N/A]") == 0.0

    def test_parse_float_empty(self):
        assert _parse_float("") == 0.0


# --- poll_gpus ---

NVIDIA_SMI_OUTPUT = """\
0, GPU-abc-123, NVIDIA A100-SXM4-80GB, 81920 MiB, 1234 MiB, 80686 MiB, 35 %, 42 C, 72.00 W, 400.00 W, 4
1, GPU-def-456, NVIDIA A100-SXM4-80GB, 81920 MiB, 0 MiB, 81920 MiB, 0 %, 38 C, 55.00 W, 400.00 W, 4
"""


class TestPollGpus:
    def test_poll_gpus_success(self):
        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = NVIDIA_SMI_OUTPUT

        with patch("radix_edge.observer.gpu_observer.subprocess.run", return_value=mock_result):
            gpus = poll_gpus()

        assert len(gpus) == 2
        assert gpus[0]["gpu_index"] == 0
        assert gpus[0]["uuid"] == "GPU-abc-123"
        assert gpus[0]["name"] == "NVIDIA A100-SXM4-80GB"
        assert gpus[0]["memory_total_mb"] == 81920
        assert gpus[0]["memory_used_mb"] == 1234
        assert gpus[0]["utilization_pct"] == 35.0
        assert gpus[0]["temperature_c"] == 42.0
        assert gpus[0]["power_draw_w"] == 72.0
        assert gpus[0]["power_limit_w"] == 400.0
        assert gpus[0]["pcie_gen"] == 4

        assert gpus[1]["gpu_index"] == 1
        assert gpus[1]["utilization_pct"] == 0.0

    def test_poll_gpus_not_found(self):
        with patch("radix_edge.observer.gpu_observer.subprocess.run", side_effect=FileNotFoundError):
            gpus = poll_gpus()
        assert gpus == []

    def test_poll_gpus_timeout(self):
        import subprocess
        with patch("radix_edge.observer.gpu_observer.subprocess.run", side_effect=subprocess.TimeoutExpired("nvidia-smi", 10)):
            gpus = poll_gpus()
        assert gpus == []

    def test_poll_gpus_bad_returncode(self):
        mock_result = MagicMock()
        mock_result.returncode = 1
        mock_result.stderr = "error"

        with patch("radix_edge.observer.gpu_observer.subprocess.run", return_value=mock_result):
            gpus = poll_gpus()
        assert gpus == []


# --- store_gpu_state ---

class TestStoreGpuState:
    def test_store_and_retrieve(self, db):
        gpus = [
            {
                "gpu_index": 0, "uuid": "GPU-abc", "name": "A100",
                "memory_total_mb": 81920, "memory_used_mb": 1000, "memory_free_mb": 80920,
                "utilization_pct": 50.0, "temperature_c": 42.0,
                "power_draw_w": 100.0, "power_limit_w": 400.0, "pcie_gen": 4,
            },
            {
                "gpu_index": 1, "uuid": "GPU-def", "name": "A100",
                "memory_total_mb": 81920, "memory_used_mb": 0, "memory_free_mb": 81920,
                "utilization_pct": 0.0, "temperature_c": 38.0,
                "power_draw_w": 55.0, "power_limit_w": 400.0, "pcie_gen": 4,
            },
        ]
        topology = {0: [1], 1: [0]}

        store_gpu_state(db, gpus, topology)

        rows = db.execute("SELECT * FROM gpus ORDER BY gpu_index").fetchall()
        assert len(rows) == 2
        assert rows[0]["name"] == "A100"
        assert rows[0]["utilization_pct"] == 50.0
        assert json.loads(rows[0]["nvlink_peers"]) == [1]
        assert rows[1]["utilization_pct"] == 0.0

        # Check observations were appended
        obs = db.execute("SELECT * FROM gpu_observations").fetchall()
        assert len(obs) == 2

    def test_store_upsert_preserves_assignment(self, db):
        """Upsert should preserve assigned_job_id from existing row."""
        gpu = {
            "gpu_index": 0, "uuid": "GPU-abc", "name": "A100",
            "memory_total_mb": 81920, "memory_used_mb": 1000, "memory_free_mb": 80920,
            "utilization_pct": 50.0, "temperature_c": 42.0,
            "power_draw_w": 100.0, "power_limit_w": 400.0, "pcie_gen": 4,
        }

        store_gpu_state(db, [gpu])

        # Manually assign a job
        db.execute("UPDATE gpus SET assigned_job_id = 'job-123', status = 'allocated' WHERE gpu_index = 0")
        db.commit()

        # Upsert again with new metrics
        gpu["utilization_pct"] = 75.0
        store_gpu_state(db, [gpu])

        row = db.execute("SELECT * FROM gpus WHERE gpu_index = 0").fetchone()
        assert row["assigned_job_id"] == "job-123"
        assert row["status"] == "allocated"
        assert row["utilization_pct"] == 75.0

    def test_store_no_topology(self, db):
        gpu = {
            "gpu_index": 0, "uuid": "GPU-abc", "name": "A100",
            "memory_total_mb": 81920, "memory_used_mb": 0, "memory_free_mb": 81920,
            "utilization_pct": 0.0, "temperature_c": 35.0,
            "power_draw_w": 50.0, "power_limit_w": 400.0, "pcie_gen": 4,
        }

        store_gpu_state(db, [gpu], topology=None)

        row = db.execute("SELECT * FROM gpus WHERE gpu_index = 0").fetchone()
        assert json.loads(row["nvlink_peers"]) == []


# --- prune_old_observations ---

class TestPruneObservations:
    def test_prune_old(self, db):
        # Insert old and new observations
        db.execute(
            "INSERT INTO gpu_observations (gpu_index, utilization_pct, memory_used_mb, temperature_c, power_draw_w, observed_at) VALUES (?, ?, ?, ?, ?, ?)",
            (0, 50.0, 1000, 42.0, 100.0, "2020-01-01T00:00:00.000000Z"),
        )
        db.execute(
            "INSERT INTO gpu_observations (gpu_index, utilization_pct, memory_used_mb, temperature_c, power_draw_w, observed_at) VALUES (?, ?, ?, ?, ?, ?)",
            (0, 60.0, 2000, 45.0, 110.0, "2099-01-01T00:00:00.000000Z"),
        )
        db.commit()

        prune_old_observations(db, retention_hours=24)

        rows = db.execute("SELECT * FROM gpu_observations").fetchall()
        assert len(rows) == 1
        assert rows[0]["utilization_pct"] == 60.0


# --- NVLink topology parsing ---

TOPO_OUTPUT_4GPU_NVLINK = """\
        GPU0    GPU1    GPU2    GPU3    CPU Affinity    NUMA Affinity
GPU0     X      NV4     NV4     NV4     0-23            0
GPU1    NV4      X      NV4     NV4     0-23            0
GPU2    NV4     NV4      X      NV4     24-47           1
GPU3    NV4     NV4     NV4      X      24-47           1
"""

TOPO_OUTPUT_PCIE_ONLY = """\
        GPU0    GPU1    CPU Affinity    NUMA Affinity
GPU0     X      PHB     0-15            0
GPU1    PHB      X      0-15            0
"""

TOPO_OUTPUT_MIXED = """\
        GPU0    GPU1    GPU2    GPU3    CPU Affinity    NUMA Affinity
GPU0     X      NV4     PHB     SYS     0-23            0
GPU1    NV4      X      SYS     PHB     0-23            0
GPU2    PHB     SYS      X      NV4     24-47           1
GPU3    SYS     PHB     NV4      X      24-47           1
"""


class TestTopologyParsing:
    def test_full_nvlink(self):
        result = parse_topology_output(TOPO_OUTPUT_4GPU_NVLINK)
        assert result is not None
        assert result[0] == [1, 2, 3]
        assert result[1] == [0, 2, 3]
        assert result[2] == [0, 1, 3]
        assert result[3] == [0, 1, 2]

    def test_pcie_only(self):
        result = parse_topology_output(TOPO_OUTPUT_PCIE_ONLY)
        assert result is None  # No NVLink connections

    def test_mixed_topology(self):
        result = parse_topology_output(TOPO_OUTPUT_MIXED)
        assert result is not None
        assert result[0] == [1]
        assert result[1] == [0]
        assert result[2] == [3]
        assert result[3] == [2]

    def test_empty_output(self):
        assert parse_topology_output("") is None

    def test_garbage_output(self):
        assert parse_topology_output("some random text\nwithout gpu info\n") is None


# --- compute_observation_vector ---

class TestObservationVector:
    def _insert_gpus(self, db, n=2, utilization=50.0, memory_used=10000, memory_total=81920,
                     temp=42.0, power_draw=100.0, power_limit=400.0, status="idle"):
        for i in range(n):
            db.execute(
                """INSERT INTO gpus (gpu_index, name, memory_total_mb, memory_used_mb, utilization_pct,
                    temperature_c, power_draw_w, power_limit_w, nvlink_peers, status)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (i, "A100", memory_total, memory_used, utilization, temp, power_draw, power_limit, "[]", status),
            )
        db.commit()

    def test_empty_cluster(self, db):
        vec = compute_observation_vector(db)
        assert vec["gpu_util_mean"] == 0.0
        assert vec["queue_depth"] == 0.0
        assert vec["memory_pressure"] == 0.0
        assert len(vec) == 11

    def test_with_gpus(self, db):
        self._insert_gpus(db, n=2, utilization=50.0, memory_used=40960, memory_total=81920,
                          temp=45.0, power_draw=200.0, power_limit=400.0)

        vec = compute_observation_vector(db)
        assert vec["gpu_util_mean"] == 50.0
        assert vec["gpu_util_std"] == 0.0  # both same utilization
        assert vec["memory_pressure"] == 0.5  # 40960/81920
        assert vec["thermal_headroom"] == 45.0 / 90.0
        assert vec["power_headroom"] == 200.0 / 400.0

    def test_with_queued_jobs(self, db):
        self._insert_gpus(db, n=1)
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j1', 'u1', 'img', 'queued')")
        db.execute("INSERT INTO jobs (job_id, user_id, image, status) VALUES ('j2', 'u1', 'img', 'queued')")
        db.commit()

        vec = compute_observation_vector(db)
        assert vec["queue_depth"] == 2.0

    def test_vector_all_keys_present(self, db):
        self._insert_gpus(db, n=1)
        vec = compute_observation_vector(db)
        expected_keys = {
            "gpu_util_mean", "gpu_util_std", "queue_depth", "queue_depth_max",
            "avg_wait_seconds", "error_rate", "throughput_jobs_per_min",
            "memory_pressure", "thermal_headroom", "power_headroom", "topology_utilization",
        }
        assert set(vec.keys()) == expected_keys


# --- GPU API endpoints ---

class TestGpuEndpoints:
    def test_list_gpus_empty(self, client, db):
        resp = client.get("/v1/gpus")
        assert resp.status_code == 200
        assert resp.json()["gpus"] == []

    def test_list_gpus_with_data(self, client, db):
        db.execute(
            "INSERT INTO gpus (gpu_index, name, memory_total_mb, status) VALUES (?, ?, ?, ?)",
            (0, "A100", 81920, "idle"),
        )
        db.execute(
            "INSERT INTO gpus (gpu_index, name, memory_total_mb, status) VALUES (?, ?, ?, ?)",
            (1, "A100", 81920, "idle"),
        )
        db.commit()

        resp = client.get("/v1/gpus")
        assert resp.status_code == 200
        gpus = resp.json()["gpus"]
        assert len(gpus) == 2
        assert gpus[0]["gpu_index"] == 0
        assert gpus[1]["gpu_index"] == 1

    def test_get_gpu_found(self, client, db):
        db.execute(
            "INSERT INTO gpus (gpu_index, name, memory_total_mb, utilization_pct, status) VALUES (?, ?, ?, ?, ?)",
            (0, "A100", 81920, 42.5, "idle"),
        )
        db.commit()

        resp = client.get("/v1/gpus/0")
        assert resp.status_code == 200
        data = resp.json()
        assert data["gpu_index"] == 0
        assert data["utilization_pct"] == 42.5

    def test_get_gpu_not_found(self, client, db):
        resp = client.get("/v1/gpus/99")
        assert resp.status_code == 404
        assert "not found" in resp.json()["detail"]
