"""GPU allocator — manages CUDA_VISIBLE_DEVICES assignments for job isolation."""

from __future__ import annotations

import logging
import threading
from typing import Optional

from radix_edge.db import get_db

logger = logging.getLogger("radix_edge.executor")

# Thread-safe lock for allocation operations
_alloc_lock = threading.Lock()


def get_available_gpus(db_conn=None) -> list[int]:
    """Return list of GPU indices that are idle (not assigned to any job)."""
    if db_conn is not None:
        rows = db_conn.execute(
            "SELECT gpu_index FROM gpus WHERE status = 'idle' ORDER BY gpu_index"
        ).fetchall()
        return [r["gpu_index"] for r in rows]

    with get_db() as db:
        rows = db.execute(
            "SELECT gpu_index FROM gpus WHERE status = 'idle' ORDER BY gpu_index"
        ).fetchall()
        return [r["gpu_index"] for r in rows]


def get_all_gpus(db_conn=None) -> list[dict]:
    """Return all GPU records."""
    if db_conn is not None:
        rows = db_conn.execute("SELECT * FROM gpus ORDER BY gpu_index").fetchall()
        return [dict(r) for r in rows]

    with get_db() as db:
        rows = db.execute("SELECT * FROM gpus ORDER BY gpu_index").fetchall()
        return [dict(r) for r in rows]


def _find_nvlink_group(gpus: list[dict], num_gpus: int, available: set[int]) -> Optional[list[int]]:
    """Find a group of NVLink-connected GPUs from the available set.

    Prefers allocating GPUs that are NVLink-connected to each other for
    better multi-GPU performance (NVLink >> PCIe bandwidth).
    """
    import json

    if num_gpus == 1:
        return None  # No topology preference for single GPU

    # Build adjacency from available GPUs
    adjacency: dict[int, set[int]] = {}
    for gpu in gpus:
        idx = gpu["gpu_index"]
        if idx not in available:
            continue
        peers_raw = gpu.get("nvlink_peers", "[]")
        peers = json.loads(peers_raw) if isinstance(peers_raw, str) else (peers_raw or [])
        adjacency[idx] = {p for p in peers if p in available}

    # Greedy: start from the GPU with most NVLink peers and expand
    for start_idx in sorted(adjacency, key=lambda i: len(adjacency.get(i, set())), reverse=True):
        group = {start_idx}
        candidates = adjacency.get(start_idx, set())
        while len(group) < num_gpus and candidates:
            # Pick candidate with most connections to current group
            best = max(candidates, key=lambda c: len(adjacency.get(c, set()) & group))
            group.add(best)
            candidates = (candidates | adjacency.get(best, set())) - group
        if len(group) >= num_gpus:
            return sorted(group)[:num_gpus]

    return None  # Couldn't find a fully connected NVLink group


def allocate(job_id: str, num_gpus: int, preferred_type: Optional[str] = None) -> Optional[list[int]]:
    """Allocate GPUs for a job. Returns list of GPU indices or None if insufficient GPUs.

    Thread-safe. Uses NVLink topology for multi-GPU allocation when possible.
    """
    with _alloc_lock:
        with get_db() as db:
            all_gpus = [dict(r) for r in db.execute("SELECT * FROM gpus ORDER BY gpu_index").fetchall()]
            available = {g["gpu_index"] for g in all_gpus if g["status"] == "idle"}

            if len(available) < num_gpus:
                return None

            # Try topology-aware allocation for multi-GPU
            indices = None
            if num_gpus > 1:
                indices = _find_nvlink_group(all_gpus, num_gpus, available)

            # Fallback: pick lowest-indexed available GPUs
            if indices is None:
                indices = sorted(available)[:num_gpus]

            # Mark GPUs as allocated
            for idx in indices:
                db.execute(
                    "UPDATE gpus SET status = 'allocated', assigned_job_id = ? WHERE gpu_index = ?",
                    (job_id, idx),
                )

            return indices


def release(job_id: str) -> list[int]:
    """Release GPUs assigned to a job. Returns the freed GPU indices."""
    with _alloc_lock:
        with get_db() as db:
            rows = db.execute(
                "SELECT gpu_index FROM gpus WHERE assigned_job_id = ?", (job_id,)
            ).fetchall()
            freed = [r["gpu_index"] for r in rows]

            if freed:
                db.execute(
                    "UPDATE gpus SET status = 'idle', assigned_job_id = NULL WHERE assigned_job_id = ?",
                    (job_id,),
                )

            return freed
