#!/usr/bin/env python3
"""Shared process metric helpers for UI panels."""

from typing import Any, Dict, List, Sequence, Tuple, Union

from ...core.gpu.base import GPUInfo

AcceleratorInfo = Union[GPUInfo, Tuple[GPUInfo, str]]


def _safe_float(value: Any) -> float:
    """Convert arbitrary values to float with safe fallback."""
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.0


def _parse_percent(value: Any) -> float:
    """Parse percentage-like values from strings or numeric types."""
    if isinstance(value, str):
        stripped = value.strip()
        if stripped.endswith("%"):
            stripped = stripped[:-1]
        value = stripped
    return _safe_float(value)


def _normalize_memory_mb(memory_value: Any, gpu_memory_total_mb: float) -> float:
    """Normalize process memory value to MB.

    Most backends already provide MB, but some AMD paths can return bytes.
    """
    value = _safe_float(memory_value)
    if value <= 0:
        return 0.0
    if gpu_memory_total_mb > 0 and value > gpu_memory_total_mb * 2:
        return value / (1024 * 1024)
    return value


def _iter_accelerators(
    accelerators: Sequence[AcceleratorInfo],
) -> List[Tuple[GPUInfo, str]]:
    """Normalize accelerator payloads into (GPUInfo, vendor) tuples."""
    normalized: List[Tuple[GPUInfo, str]] = []
    for entry in accelerators:
        if isinstance(entry, tuple):
            gpu, vendor = entry
            normalized.append((gpu, str(vendor)))
        else:
            normalized.append((entry, "gpu"))
    return normalized


def _device_label(vendor: str, index: int) -> str:
    """Build a compact accelerator label used by process tables."""
    prefix = "NPU" if vendor.endswith("_npu") else "GPU"
    return f"{prefix}{index}"


def _calculate_device_metrics(
    process: Dict[str, Any],
    accelerator: GPUInfo,
) -> Tuple[bool, float, float]:
    """Calculate process utilization and memory-percent for one accelerator."""
    pid = process.get("pid")
    if pid is None:
        return False, 0.0, 0.0

    total_memory_mb = _safe_float(accelerator.memory_total)
    if total_memory_mb <= 0:
        total_memory_mb = 0.0

    normalized_total_proc_memory = sum(
        _normalize_memory_mb(item.get("memory", 0.0), total_memory_mb)
        for item in accelerator.processes
    )

    found_match = False
    max_util = 0.0
    max_memory_percent = 0.0

    for accelerator_proc in accelerator.processes:
        if accelerator_proc.get("pid") != pid:
            continue

        found_match = True
        proc_memory_mb = _normalize_memory_mb(
            accelerator_proc.get("memory", 0.0), total_memory_mb
        )

        if accelerator_proc.get("sm_util") is not None:
            current_util = _parse_percent(accelerator_proc.get("sm_util"))
        elif accelerator_proc.get("cu_occupancy") is not None:
            current_util = _parse_percent(accelerator_proc.get("cu_occupancy"))
        elif accelerator_proc.get("gpu_util") is not None:
            current_util = _parse_percent(accelerator_proc.get("gpu_util"))
        else:
            current_util = 0.0
            if normalized_total_proc_memory > 0:
                memory_ratio = proc_memory_mb / normalized_total_proc_memory
                current_util = _safe_float(accelerator.utilization) * memory_ratio

        current_memory_percent = 0.0
        if total_memory_mb > 0:
            current_memory_percent = (proc_memory_mb / total_memory_mb) * 100.0

        max_util = max(max_util, current_util)
        max_memory_percent = max(max_memory_percent, current_memory_percent)

    return found_match, max(0.0, max_util), max(0.0, max_memory_percent)


def calculate_process_device_breakdown(
    process: Dict[str, Any],
    accelerators: Sequence[AcceleratorInfo],
    max_devices: int = 3,
) -> List[Dict[str, Any]]:
    """Return matching accelerators sorted by descending utilization.

    Each entry contains:
    - ``device``: label such as ``GPU0`` or ``NPU0``
    - ``utilization``: per-process utilization percentage on that device
    - ``memory_percent``: process memory percentage on that device
    """
    if max_devices <= 0:
        return []

    matches: List[Dict[str, Any]] = []
    for accelerator, vendor in _iter_accelerators(accelerators):
        found, utilization, memory_percent = _calculate_device_metrics(
            process, accelerator
        )
        if not found:
            continue

        matches.append(
            {
                "device": _device_label(vendor, accelerator.index),
                "utilization": utilization,
                "memory_percent": memory_percent,
            }
        )

    matches.sort(
        key=lambda item: (
            float(item.get("utilization", 0.0)),
            float(item.get("memory_percent", 0.0)),
            str(item.get("device", "")),
        ),
        reverse=True,
    )
    return matches[:max_devices]


def calculate_process_gpu_metrics(
    process: Dict[str, Any], gpus: List[GPUInfo]
) -> Tuple[float, float]:
    """Calculate per-process GPU and VRAM percentages.

    Returns:
        Tuple of (gpu_percent, vram_percent), using the highest matching GPU values.
    """
    if not gpus:
        return 0.0, 0.0

    breakdown = calculate_process_device_breakdown(process, gpus, max_devices=len(gpus))
    if not breakdown:
        return 0.0, 0.0

    gpu_percent = max(float(item.get("utilization", 0.0)) for item in breakdown)
    vram_percent = max(float(item.get("memory_percent", 0.0)) for item in breakdown)
    return max(0.0, gpu_percent), max(0.0, vram_percent)
