#!/usr/bin/env python3
"""Shared helpers for process-level accelerator metrics in UI panels."""

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

from ...core.gpu.base import GPUInfo

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


def _safe_float(value: Any) -> float:
    """Convert a value to float, returning 0.0 on invalid input."""
    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 _clamp_percent(value: Any) -> float:
    """Clamp percentage-like values to the display-safe 0..100 range."""
    parsed = _parse_percent(value)
    return max(0.0, min(100.0, parsed))


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

    Most backends already provide MB, but some AMD paths return bytes, which
    we detect by comparing to total memory and convert accordingly.
    """
    value = _safe_float(memory_value)
    if value <= 0:
        return 0.0

    total_memory_mb = max(0.0, _safe_float(gpu_memory_total_mb))
    mb_candidate = value
    kib_candidate = value / 1024.0
    bytes_candidate = value / (1024.0 * 1024.0)
    candidates = [mb_candidate, kib_candidate, bytes_candidate]

    if total_memory_mb > 0:
        # Prefer the largest plausible candidate that fits within device bounds.
        upper_bound = max(total_memory_mb * 2.0, total_memory_mb + 1024.0)
        plausible = [candidate for candidate in candidates if candidate <= upper_bound]
        if plausible:
            return max(0.0, max(plausible))

    # If total memory is unknown, fall back to magnitude-based conversion.
    if value >= 1024.0 * 1024.0 * 1024.0:
        return bytes_candidate
    if value >= 1024.0 * 1024.0:
        return kib_candidate

    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 _is_npu_vendor(vendor: str) -> bool:
    """Return True when vendor string represents an NPU class accelerator."""
    return str(vendor).endswith("_npu")


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


def build_accelerator_identities(
    accelerators: Sequence[AcceleratorInfo],
) -> List[Dict[str, Any]]:
    """Build stable per-class accelerator identities for UI display/grouping.

    The global index is scoped per accelerator class:
    - GPUs are numbered ``GPU0..GPU{n}`` across all GPU vendors.
    - NPUs are numbered ``NPU0..NPU{n}`` across all NPU vendors.
    """
    normalized = _iter_accelerators(accelerators)
    identities: List[Dict[str, Any]] = []
    gpu_counter = 0
    npu_counter = 0

    for accelerator, vendor in normalized:
        local_index = int(accelerator.index)
        local_label = _device_label(vendor, local_index)
        if _is_npu_vendor(vendor):
            global_index = npu_counter
            npu_counter += 1
            device_class = "npu"
        else:
            global_index = gpu_counter
            gpu_counter += 1
            device_class = "gpu"
        global_label = _device_label(vendor, global_index)

        identities.append(
            {
                "vendor": vendor,
                "device_class": device_class,
                "local_index": local_index,
                "local_label": local_label,
                "global_index": global_index,
                "global_label": global_label,
            }
        )

    return identities


def _iter_accelerators_with_identity(
    accelerators: Sequence[AcceleratorInfo],
) -> List[Tuple[GPUInfo, str, Dict[str, Any]]]:
    """Normalize accelerators and pair each with per-class identity metadata."""
    normalized = _iter_accelerators(accelerators)
    identities = build_accelerator_identities(normalized)
    return [
        (accelerator, vendor, identity)
        for (accelerator, vendor), identity in zip(normalized, identities)
    ]


def _calculate_device_metrics(
    process: Dict[str, Any],
    accelerator: GPUInfo,
) -> Tuple[bool, float, float, float, float]:
    """Compute per-process utilization and memory stats for one accelerator."""
    pid = process.get("pid")
    if pid is None:
        return False, 0.0, 0.0, 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
    proc_memory_mb_total = 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
        )
        proc_memory_mb_total += proc_memory_mb

        if accelerator_proc.get("sm_util") is not None:
            current_util = _clamp_percent(accelerator_proc.get("sm_util"))
        elif accelerator_proc.get("cu_occupancy") is not None:
            current_util = _clamp_percent(accelerator_proc.get("cu_occupancy"))
        elif accelerator_proc.get("gpu_util") is not None:
            current_util = _clamp_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 = _clamp_percent(
                    _safe_float(accelerator.utilization) * memory_ratio
                )

        max_util = max(max_util, _clamp_percent(current_util))

    memory_percent = 0.0
    if total_memory_mb > 0:
        memory_percent = _clamp_percent(
            (proc_memory_mb_total / total_memory_mb) * 100.0
        )

    return (
        found_match,
        _clamp_percent(max_util),
        _clamp_percent(memory_percent),
        max(0.0, proc_memory_mb_total),
        max(0.0, total_memory_mb),
    )


def recommended_device_display_limit(total_accelerators: int) -> int:
    """Return a device-count cap that stays readable for small/large systems."""
    if total_accelerators <= 0:
        return 1
    if total_accelerators <= 3:
        return total_accelerators
    if total_accelerators <= 8:
        return 3
    return 2


def _format_memory_short(memory_mb: float) -> str:
    """Format MB into compact MB/GB text."""
    memory_mb = max(0.0, _safe_float(memory_mb))
    if memory_mb >= 1024:
        return f"{memory_mb / 1024:.1f}G"
    return f"{memory_mb:.0f}M"


def format_device_breakdown(
    breakdown: Sequence[Mapping[str, Any]],
    max_devices: int,
    include_utilization: bool = True,
    include_memory: bool = False,
) -> str:
    """Format per-process accelerator breakdown for compact UI columns."""
    if not breakdown or max_devices <= 0:
        return ""

    limited = list(breakdown[:max_devices])
    parts: List[str] = []
    for entry in limited:
        device = str(entry.get("device", "")).strip()
        if not device:
            continue

        util = _clamp_percent(entry.get("utilization", 0.0))
        memory_mb = float(entry.get("memory_mb", 0.0))
        if include_utilization and include_memory:
            parts.append(f"{device}:{util:.0f}%/{_format_memory_short(memory_mb)}")
        elif include_utilization:
            parts.append(f"{device}:{util:.0f}%")
        elif include_memory:
            parts.append(f"{device}:{_format_memory_short(memory_mb)}")
        else:
            parts.append(device)

    remaining = max(0, len(breakdown) - len(limited))
    if remaining:
        parts.append(f"+{remaining}")
    return ",".join(parts)


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``: globally unique per-class label such as ``GPU0`` or ``NPU0``
    - ``device_global_label``/``device_global_index``: explicit global id fields
    - ``device_local_label``/``index``: original backend-local device id fields
    - ``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, identity in _iter_accelerators_with_identity(accelerators):
        (
            found,
            utilization,
            memory_percent,
            memory_mb,
            memory_total_mb,
        ) = _calculate_device_metrics(
            process,
            accelerator,
        )
        if not found:
            continue

        matches.append(
            {
                "device": str(identity["global_label"]),
                "device_global_label": str(identity["global_label"]),
                "device_global_index": int(identity["global_index"]),
                "device_local_label": str(identity["local_label"]),
                "vendor": vendor,
                "index": accelerator.index,
                "local_index": accelerator.index,
                "utilization": utilization,
                "memory_percent": memory_percent,
                "memory_mb": memory_mb,
                "memory_total_mb": memory_total_mb,
            }
        )

    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 summarize_process_device_resources(
    processes: Sequence[Dict[str, Any]],
    accelerators: Sequence[AcceleratorInfo],
) -> List[Dict[str, Any]]:
    """Aggregate per-device usage for a process cohort."""
    normalized = _iter_accelerators(accelerators)
    if not normalized or not processes:
        return []

    max_devices = max(1, len(normalized))
    device_totals: Dict[str, Dict[str, Any]] = {}
    for process in processes:
        for item in calculate_process_device_breakdown(
            process, normalized, max_devices=max_devices
        ):
            device = str(
                item.get("device_global_label")
                or item.get("device")
                or item.get("device_local_label")
                or ""
            )
            if not device:
                continue
            existing = device_totals.setdefault(
                device,
                {
                    "device": device,
                    "device_global_label": device,
                    "device_global_index": int(
                        item.get("device_global_index", item.get("index", 0))
                    ),
                    "device_local_label": str(
                        item.get("device_local_label", item.get("device", ""))
                    ),
                    "index": int(item.get("index", 0)),
                    "vendor": str(item.get("vendor", "")),
                    "process_count": 0,
                    "utilization": 0.0,
                    "memory_used_mb": 0.0,
                    "memory_total_mb": float(item.get("memory_total_mb", 0.0)),
                },
            )
            existing["process_count"] += 1
            existing["utilization"] = min(
                100.0,
                float(existing["utilization"]) + float(item.get("utilization", 0)),
            )
            existing["memory_used_mb"] += float(item.get("memory_mb", 0.0))
            existing["memory_total_mb"] = max(
                float(existing["memory_total_mb"]),
                float(item.get("memory_total_mb", 0.0)),
            )

    summaries = list(device_totals.values())
    summaries.sort(
        key=lambda item: (
            float(item.get("utilization", 0.0)),
            float(item.get("memory_used_mb", 0.0)),
            int(item.get("process_count", 0)),
            str(item.get("device", "")),
        ),
        reverse=True,
    )
    return summaries


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(
        _clamp_percent(item.get("utilization", 0.0)) for item in breakdown
    )
    vram_percent = max(
        _clamp_percent(item.get("memory_percent", 0.0)) for item in breakdown
    )
    return _clamp_percent(gpu_percent), _clamp_percent(vram_percent)
