#!/usr/bin/env python3
"""GPU information panel component."""

import curses
import os
from typing import Any, Dict, List, Optional, Tuple

import psutil

from ...application.process_query import filter_processes
from ...core.gpu.base import GPUInfo
from ..display import Display
from .process_metrics import (
    build_accelerator_identities,
    calculate_process_gpu_metrics,
)


class GPUPanel:
    """Renders GPU information panel."""

    def __init__(self, display: Display):
        """Initialize the GPU panel.

        Args:
            display: Display instance
        """
        self.display = display

    def _theme_attr(self, role: str, fallback_pair: int = 0) -> int:
        """Resolve theme attribute with compatibility fallback for test doubles."""
        getter = getattr(self.display, "get_theme_attr", None)
        if callable(getter):
            return int(getter(role))
        return curses.color_pair(fallback_pair)

    @staticmethod
    def _clamp_percent(value: Any) -> float:
        """Clamp percentage-like values to a safe display range."""
        try:
            parsed = float(value)
        except (TypeError, ValueError):
            return 0.0
        return max(0.0, min(100.0, parsed))

    @staticmethod
    def _device_prefix(vendor: str) -> str:
        """Map vendor value to compact device prefix."""
        return "NPU" if str(vendor).endswith("_npu") else "GPU"

    @staticmethod
    def _memory_marker(gpu: GPUInfo) -> str:
        """Return a compact memory-model marker for display labels."""
        model = str(getattr(gpu, "memory_model", "")).strip().lower()
        if model == "shared":
            return "~"
        if model == "hybrid":
            return "+"
        return " "

    def _summarize_cmdline(self, cmdline: List[str], fallback_name: str) -> str:
        """Create a short process label from command arguments."""
        if not cmdline:
            return fallback_name

        executable = os.path.basename(cmdline[0]) if cmdline[0] else fallback_name
        lowered = executable.lower()

        if lowered.startswith("python"):
            if len(cmdline) >= 3 and cmdline[1] == "-m":
                module = cmdline[2]
                return f"python -m {module}"
            for arg in cmdline[1:]:
                if arg.endswith(".py"):
                    return f"python:{os.path.basename(arg)}"
            return executable

        if len(cmdline) > 1:
            return f"{executable} {cmdline[1]}"

        return executable or fallback_name

    def _resolve_runtime_process_fields(
        self,
        pid: int,
        fallback_name: str,
        fallback_ppid: Optional[int],
        runtime_cache: Dict[int, Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Resolve runtime process name/ppid/cmdline with per-render caching."""
        cached = runtime_cache.get(pid)
        if cached is not None:
            return cached

        resolved: Dict[str, Any] = {
            "name": fallback_name,
            "ppid": fallback_ppid,
            "cmdline_summary": "",
            "cmdline": "",
        }

        try:
            process = psutil.Process(pid)
            with process.oneshot():
                ppid = process.ppid()
                cmdline = process.cmdline()

            if ppid > 0:
                resolved["ppid"] = ppid

            if cmdline:
                resolved["cmdline"] = " ".join(cmdline)
                resolved["cmdline_summary"] = self._summarize_cmdline(
                    cmdline, fallback_name
                )
                resolved["name"] = resolved["cmdline_summary"]
        except (
            psutil.NoSuchProcess,
            psutil.AccessDenied,
            psutil.ZombieProcess,
            OSError,
            ValueError,
        ):
            pass

        runtime_cache[pid] = resolved
        return resolved

    def _build_context_processes(
        self,
        gpu: GPUInfo,
        vendor: str,
        runtime_cache: Dict[int, Dict[str, Any]],
        device_name: Optional[str] = None,
        device_global_index: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """Attach GPU and runtime context fields used by filtering/rendering."""
        context_processes: List[Dict[str, Any]] = []
        device_prefix = self._device_prefix(vendor)
        local_device_name = f"{device_prefix}{gpu.index}"
        resolved_device_name = device_name or local_device_name
        resolved_global_index = (
            int(gpu.index)
            if device_global_index is None
            else int(device_global_index)
        )
        for process in gpu.processes:
            process_with_context = dict(process)
            process_with_context["is_gpu_process"] = True
            process_with_context["gpu_name"] = gpu.name
            process_with_context["gpu_vendor"] = vendor
            process_with_context["device_label"] = device_prefix
            process_with_context["device_index"] = gpu.index
            process_with_context["device_local_label"] = local_device_name
            process_with_context["device_global_index"] = resolved_global_index
            process_with_context["device_global_label"] = resolved_device_name
            process_with_context["device"] = resolved_device_name

            pid = process_with_context.get("pid")
            if isinstance(pid, int):
                fallback_name = str(process_with_context.get("name", ""))
                ppid_value = process_with_context.get("ppid")
                fallback_ppid = ppid_value if isinstance(ppid_value, int) else None
                runtime = self._resolve_runtime_process_fields(
                    pid, fallback_name, fallback_ppid, runtime_cache
                )
                process_with_context["name"] = runtime["name"]
                if runtime["ppid"] is not None:
                    process_with_context["ppid"] = runtime["ppid"]
                if runtime["cmdline_summary"]:
                    process_with_context["cmdline_summary"] = runtime["cmdline_summary"]
                if runtime["cmdline"]:
                    process_with_context["cmdline"] = runtime["cmdline"]

                # Populate fallback per-process accelerator utilization so NPUs
                # (and other backends without process-level util counters)
                # still show useful percentages in the tree view.
                gpu_util, vram_percent = calculate_process_gpu_metrics(
                    process_with_context, [gpu]
                )
                process_with_context["gpu_util"] = gpu_util
                process_with_context["vram_percent"] = vram_percent

            context_processes.append(process_with_context)

        return context_processes

    @staticmethod
    def _build_accelerator_metadata(gpu: GPUInfo) -> str:
        """Build optional metadata line (driver/toolkit versions)."""
        parts: List[str] = []
        if gpu.driver_version:
            parts.append(f"Driver: {gpu.driver_version}")
        if gpu.cuda_version:
            parts.append(f"CUDA: {gpu.cuda_version}")
        if gpu.rocm_version:
            parts.append(f"ROCm: {gpu.rocm_version}")
        return " | ".join(parts)

    def _render_gpu_header(
        self,
        gpu: GPUInfo,
        y: int,
        x: int,
        vendor: str,
        device_label: str = "GPU",
        device_display_index: Optional[int] = None,
    ) -> int:
        """Render GPU header with name and vendor.

        Args:
            gpu: GPU information
            y: Starting Y coordinate
            x: Starting X coordinate
            vendor: GPU vendor string
            device_label: Label prefix for accelerator type

        Returns:
            Next Y coordinate
        """
        vendor_prefix = ""
        if all(v.upper() not in gpu.name.upper() for v in ["NVIDIA", "AMD", "INTEL"]):
            vendor_prefix = f"{vendor.upper()} "

        display_index = (
            int(gpu.index)
            if device_display_index is None
            else int(device_display_index)
        )
        header = f"{device_label} {display_index}: {vendor_prefix}{gpu.name}"
        header_attr = self._theme_attr("header", fallback_pair=1) | curses.A_BOLD
        self.display.safe_addstr(y, x, header, header_attr)
        metadata_line = self._build_accelerator_metadata(gpu)
        if metadata_line:
            self.display.safe_addstr(
                y + 1, x, metadata_line, self._theme_attr("muted", fallback_pair=5)
            )
        return y + 2

    def _render_utilization_bar(
        self, gpu: GPUInfo, y: int, x: int, width: int, device_label: str = "GPU"
    ) -> int:
        """Render GPU utilization bar.

        Args:
            gpu: GPU information
            y: Starting Y coordinate
            x: Starting X coordinate
            width: Maximum width for the bar
            device_label: Label prefix for accelerator type

        Returns:
            Next Y coordinate
        """
        utilization = self._clamp_percent(gpu.utilization)
        bar, color = self.display.create_bar(utilization, width)
        line = f"{device_label[:4]:<4} [{bar}] {utilization:5.1f}%"
        self.display.safe_addstr(y, x, line, color)
        return y + 1

    def _render_memory_bar(self, gpu: GPUInfo, y: int, x: int, width: int) -> int:
        """Render memory usage bar.

        Args:
            gpu: GPU information (with memory_used and memory_total in MB)
            y: Starting Y coordinate
            x: Starting X coordinate
            width: Maximum width for the bar

        Returns:
            Next Y coordinate

        Note:
            The GPU memory values in the GPUInfo object are in MB (megabytes).
            These values are converted to GB (gigabytes) for display.
        """
        known_total = gpu.memory_total > 0
        mem_percent = (gpu.memory_used / gpu.memory_total * 100) if known_total else 0
        mem_percent = self._clamp_percent(mem_percent)
        bar, color = self.display.create_bar(mem_percent, width)
        marker = self._memory_marker(gpu)
        has_theme_attr = callable(getattr(self.display, "get_theme_attr", None))
        if marker != " " and has_theme_attr:
            color = self._theme_attr("warning", fallback_pair=3) | curses.A_BOLD
        elif not known_total and has_theme_attr:
            color = self._theme_attr("muted", fallback_pair=5)

        # Convert memory values from MB to GB for display
        mem_used_gb = gpu.memory_used / 1024
        mem_total_gb = gpu.memory_total / 1024
        if known_total:
            line = (
                f"Mem{marker} [{bar}] {mem_used_gb:5.1f}GB / "
                f"{mem_total_gb:5.1f}GB"
            )
        else:
            line = f"Mem{marker} [{bar}] {mem_used_gb:5.1f}GB /   N/A"
        self.display.safe_addstr(y, x, line, color)
        return y + 1

    def _render_temperature(self, gpu: GPUInfo, y: int, x: int) -> int:
        """Render temperature information.

        Args:
            gpu: GPU information
            y: Starting Y coordinate
            x: Starting X coordinate

        Returns:
            Next Y coordinate
        """
        if gpu.temperature > 0:
            line = f"Temp  {gpu.temperature:5.1f}°C"
            self.display.safe_addstr(
                y, x, line, self.display.get_color(gpu.temperature)
            )
            return y + 1
        return y

    def _render_power(self, gpu: GPUInfo, y: int, x: int) -> int:
        """Render power usage information.

        Args:
            gpu: GPU information
            y: Starting Y coordinate
            x: Starting X coordinate

        Returns:
            Next Y coordinate
        """
        if gpu.power_draw > 0 and gpu.power_limit > 0:
            power_percent = gpu.power_draw / gpu.power_limit * 100
            line = f"Power {gpu.power_draw:5.1f}W / {gpu.power_limit:5.1f}W"
            self.display.safe_addstr(y, x, line, self.display.get_color(power_percent))
            return y + 1
        return y

    def _build_process_tree(
        self, processes: List[Dict[str, Any]]
    ) -> Tuple[List[Dict[str, Any]], Dict[int, List[Dict[str, Any]]]]:
        """Build process tree structure from flat process list.

        Args:
            processes: List of process dictionaries with pid and ppid

        Returns:
            Tuple of (root_processes, children_map) where:
            - root_processes: Processes with no parent in the GPU process list
            - children_map: Dict mapping parent PID to list of child processes
        """
        # Build PID to process mapping
        pid_to_proc = {proc["pid"]: proc for proc in processes}

        # Build children mapping
        children_map: Dict[int, List[Dict[str, Any]]] = {}
        root_processes: List[Dict[str, Any]] = []

        for proc in processes:
            ppid = proc.get("ppid")
            if ppid and ppid in pid_to_proc:
                # Parent is in GPU process list
                if ppid not in children_map:
                    children_map[ppid] = []
                children_map[ppid].append(proc)
            else:
                # Root process (no parent in GPU list)
                root_processes.append(proc)

        return root_processes, children_map

    def _render_processes(
        self,
        gpu: GPUInfo,
        vendor: str,
        y: int,
        x: int,
        selected_pid: Optional[int] = None,
        process_filter: str = "",
        device_label: str = "GPU",
        device_name: Optional[str] = None,
        device_display_index: Optional[int] = None,
    ) -> int:
        """Render GPU process information with tree hierarchy.

        Args:
            gpu: GPU information
            y: Starting Y coordinate
            x: Starting X coordinate
            selected_pid: PID of selected process for highlighting
            process_filter: Optional process name/PID filter query

        Returns:
            Next Y coordinate
        """
        if not gpu.processes:
            return y

        runtime_cache: Dict[int, Dict[str, Any]] = {}
        context_processes = self._build_context_processes(
            gpu,
            vendor,
            runtime_cache,
            device_name=device_name,
            device_global_index=device_display_index,
        )

        filtered_processes = filter_processes(context_processes, process_filter)
        if not filtered_processes:
            return y

        display_device_name = device_name
        if not display_device_name:
            if device_display_index is None:
                display_device_name = f"{device_label}{gpu.index}"
            else:
                display_device_name = f"{device_label}{int(device_display_index)}"

        y += 1
        self.display.safe_addstr(
            y,
            x,
            f"Running Processes on {display_device_name}:",
            self._theme_attr("chip", fallback_pair=5),
        )
        y += 1

        # Build process tree
        root_processes, children_map = self._build_process_tree(filtered_processes)

        # Render processes with tree structure
        for i, proc in enumerate(root_processes):
            is_last_root = i == len(root_processes) - 1
            y = self._render_process_tree_node(
                proc,
                y,
                x,
                selected_pid,
                children_map,
                is_last_root,
                "",
                device_label,
            )

        return y + 1

    def _render_process_tree_node(
        self,
        proc: Dict[str, Any],
        y: int,
        x: int,
        selected_pid: Optional[int],
        children_map: Dict[int, List[Dict[str, Any]]],
        is_last: bool,
        prefix: str,
        device_label: str = "GPU",
    ) -> int:
        """Recursively render a process tree node and its children.

        Args:
            proc: Process dictionary
            y: Current Y coordinate
            x: Starting X coordinate
            selected_pid: PID of selected process
            children_map: Map of parent PID to children
            is_last: Whether this is the last sibling
            prefix: Current tree prefix string

        Returns:
            Next Y coordinate
        """
        # Determine tree symbols
        if prefix:
            tree_symbol = "└── " if is_last else "├── "
        else:
            tree_symbol = "├── " if not is_last else "└── "

        # Render per-process device id explicitly for multi-device readability.
        device_tag = str(proc.get("device", f"{device_label}?"))
        util_text = "--.-"
        if "sm_util" in proc:
            gpu_util = self._clamp_percent(proc["sm_util"])
            util_text = f"{gpu_util:4.1f}"
        elif "cu_occupancy" in proc and proc["cu_occupancy"] is not None:
            gpu_util = self._clamp_percent(proc["cu_occupancy"])
            util_text = f"{gpu_util:4.1f}"
        elif "gpu_util" in proc and proc["gpu_util"] is not None:
            gpu_util = self._clamp_percent(proc["gpu_util"])
            util_text = f"{gpu_util:4.1f}"
        gpu_util_str = f"{device_tag} {util_text}% | "

        # Process memory might be in bytes instead of MB
        # for some GPU implementations
        memory_value = proc.get("memory", 0.0)
        try:
            memory_value = float(memory_value)
        except (TypeError, ValueError):
            memory_value = 0.0

        # If value is extremely large (over 1000 GB when interpreted as MB),
        # assume it's in bytes and convert directly to GB
        if memory_value > 1000 * 1024:  # If > 1000 GB as MB
            # Convert bytes to GB
            mem_gb = memory_value / (1024 * 1024 * 1024)
        else:
            # Otherwise, assume it's in MB as documented and convert to GB
            mem_gb = memory_value / 1024

        # Format with appropriate precision based on size
        if mem_gb < 0.01:  # Very tiny values (under 10MB)
            mem_format = f"{mem_gb*1024:.1f}MB"
        elif mem_gb < 0.1:  # Small values (under 100MB)
            mem_format = f"{mem_gb:.3f}GB"
        elif mem_gb < 1:  # Medium values (under 1GB)
            mem_format = f"{mem_gb:.2f}GB"
        else:  # Large values (1GB and above)
            mem_format = f"{mem_gb:.1f}GB"

        # Build the display line with tree symbols
        line = (
            f"{prefix}{tree_symbol}PID {proc['pid']:6d} | {gpu_util_str}"
            f"{mem_format:>9s} | {proc['name']}"
        )

        # Apply selection highlighting
        is_selected = selected_pid is not None and proc["pid"] == selected_pid
        attr = (
            self._theme_attr("selected", fallback_pair=2)
            if is_selected
            else self._theme_attr("info", fallback_pair=6)
        )
        self.display.safe_addstr(y, x, line, attr)
        y += 1

        # Render children recursively
        children = children_map.get(proc["pid"], [])
        for i, child in enumerate(children):
            is_last_child = i == len(children) - 1
            # Update prefix for children
            child_prefix = prefix + ("    " if is_last else "│   ")
            y = self._render_process_tree_node(
                child,
                y,
                x,
                selected_pid,
                children_map,
                is_last_child,
                child_prefix,
                device_label,
            )

        return y

    def render(
        self,
        gpu_info: List[tuple[GPUInfo, str]],
        start_y: int = 3,
        indent: int = 2,
        selected_pid: Optional[int] = None,
        process_filter: str = "",
        device_label: str = "GPU",
    ) -> int:
        """Render the complete GPU panel.

        Args:
            gpu_info: List of tuples containing (GPU information, vendor string)
            start_y: Starting Y coordinate
            indent: Left indentation
            selected_pid: PID of selected process for highlighting
            process_filter: Optional process name/PID filter query
            device_label: Label prefix for accelerator type

        Returns:
            Next Y coordinate
        """
        if not gpu_info:
            self.display.safe_addstr(
                start_y,
                indent,
                f"No compatible {device_label}s detected",
                self._theme_attr("warning", fallback_pair=3),
            )
            return start_y + 1

        y = start_y
        bar_width = min(50, self.display.width - 35)
        identities = build_accelerator_identities(gpu_info)

        for (gpu, vendor), identity in zip(gpu_info, identities):
            device_global_index = int(identity.get("global_index", gpu.index))
            device_global_label = str(
                identity.get(
                    "global_label",
                    f"{self._device_prefix(vendor)}{device_global_index}",
                )
            )

            # Render GPU sections
            y = self._render_gpu_header(
                gpu,
                y,
                indent,
                vendor,
                device_label,
                device_display_index=device_global_index,
            )
            y = self._render_utilization_bar(gpu, y, indent, bar_width, device_label)
            y = self._render_memory_bar(gpu, y, indent, bar_width)
            y = self._render_temperature(gpu, y, indent)
            y = self._render_power(gpu, y, indent)
            y = self._render_processes(
                gpu,
                vendor,
                y,
                indent,
                selected_pid,
                process_filter,
                device_label,
                device_name=device_global_label,
                device_display_index=device_global_index,
            )

        return y

    def get_all_gpu_processes(
        self, gpu_info: List[tuple[GPUInfo, str]], process_filter: str = ""
    ) -> List[Dict[str, Any]]:
        """Get all processes across all GPUs as a flat list.

        Args:
            gpu_info: List of tuples containing (GPU information, vendor string)
            process_filter: Optional process name/PID filter query

        Returns:
            List of process dictionaries with PID, name, and memory
        """
        all_processes = []
        seen_pids = set()
        runtime_cache: Dict[int, Dict[str, Any]] = {}
        identities = build_accelerator_identities(gpu_info)

        for (gpu, vendor), identity in zip(gpu_info, identities):
            local_device_name = f"{self._device_prefix(vendor)}{gpu.index}"
            global_device_name = str(
                identity.get("global_label", local_device_name)
            )
            global_device_index = int(identity.get("global_index", gpu.index))
            context_processes = self._build_context_processes(
                gpu,
                vendor,
                runtime_cache,
                device_name=global_device_name,
                device_global_index=global_device_index,
            )

            filtered_processes = filter_processes(context_processes, process_filter)
            for proc in filtered_processes:
                pid = proc["pid"]
                # Avoid duplicates (process can be on multiple GPUs)
                if pid not in seen_pids:
                    seen_pids.add(pid)
                    all_processes.append(
                        {
                            "pid": pid,
                            "ppid": proc.get("ppid"),
                            "name": proc["name"],
                            "memory": proc["memory"],
                            "is_gpu_process": True,
                            "gpu_name": proc.get("gpu_name", gpu.name),
                            "gpu_vendor": proc.get("gpu_vendor", vendor),
                            "device": proc.get(
                                "device",
                                global_device_name,
                            ),
                            "device_index": proc.get("device_index", gpu.index),
                            "device_local_label": proc.get(
                                "device_local_label", local_device_name
                            ),
                            "device_global_index": proc.get(
                                "device_global_index", global_device_index
                            ),
                            "device_global_label": proc.get(
                                "device_global_label", global_device_name
                            ),
                            "cmdline_summary": proc.get("cmdline_summary", ""),
                            "cmdline": proc.get("cmdline", ""),
                        }
                    )

        return all_processes

    def get_selected_process_by_index(
        self,
        gpu_info: List[tuple[GPUInfo, str]],
        selected_index: int,
        process_filter: str = "",
    ) -> Optional[Dict[str, Any]]:
        """Get process by selection index.

        Args:
            gpu_info: List of tuples containing (GPU information, vendor string)
            selected_index: Index of selected process
            process_filter: Optional process name/PID filter query

        Returns:
            Process dict or None if invalid index
        """
        all_processes = self.get_all_gpu_processes(gpu_info, process_filter)

        if 0 <= selected_index < len(all_processes):
            return all_processes[selected_index]

        return None
