"""Linux GPU performance optimization for AMD (RDNA3/sysfs) and NVIDIA."""
from __future__ import annotations

import glob
import logging
from pathlib import Path
from typing import Optional

from mesh_boost.optimizer.base import (
    IS_LINUX, OptimizationResult, RollbackManager,
    is_admin, run_cmd, sysfs_read, sysfs_write,
)

logger = logging.getLogger(__name__)

# AMD GPU sysfs base
AMDGPU_HWMON = "/sys/class/drm/card*/device/hwmon/hwmon*"
AMDGPU_DEVICE = "/sys/class/drm/card*/device"


class LinuxGPUOptimizer:
    """Optimize GPU settings on Linux for AMD (sysfs) and NVIDIA (nvidia-smi)."""

    def __init__(self):
        self._rollback = RollbackManager("linux_gpu")

    def apply(self, dry_run: bool = False) -> OptimizationResult:
        result = OptimizationResult("linux_gpu")

        if not IS_LINUX:
            result.skip("all", "not Linux")
            return result

        vendor = self._detect_gpu_vendor()
        if not vendor:
            result.skip("gpu", "no supported GPU detected")
            return result

        if "amd" in vendor:
            self._apply_amd_power_profile(result, dry_run)
            self._apply_amd_performance_level(result, dry_run)
            self._apply_amd_fan_control(result, dry_run)
            self._apply_amd_clock_boost(result, dry_run)
        if "nvidia" in vendor:
            self._apply_nvidia_perf_mode(result, dry_run)
            self._apply_nvidia_power_limit(result, dry_run)

        return result

    def restore(self) -> OptimizationResult:
        result = OptimizationResult("linux_gpu_restore")
        if not IS_LINUX:
            return result

        self._restore_amd(result)
        self._restore_nvidia(result)

        self._rollback.clear()
        return result

    # ── AMD GPU (sysfs/RDNA3) ────────────────────────────────────────────

    def _apply_amd_power_profile(self, result: OptimizationResult, dry_run: bool):
        """Set AMD GPU power profile to 3D_FULL_SCREEN (gaming profile).

        Available profiles: BOOTUP_DEFAULT, 3D_FULL_SCREEN, POWER_SAVING,
        VIDEO, VR, COMPUTE, CUSTOM. 3D_FULL_SCREEN gives highest clocks
        and most responsive power state transitions.
        """
        if dry_run:
            result.skip("amd_power_profile", "dry_run")
            return

        device_paths = self._get_amd_device_paths()
        if not device_paths:
            result.skip("amd_power_profile", "no AMD GPU sysfs found")
            return

        for dev in device_paths:
            profile_path = f"{dev}/pp_power_profile_mode"
            if not Path(profile_path).exists():
                continue

            orig = sysfs_read(profile_path)
            self._rollback.save_original("amd_power_profile", orig)

            # Profile 1 = 3D_FULL_SCREEN
            if sysfs_write(profile_path, "1"):
                result.ok("amd_power_profile=3D_FULL_SCREEN")
            else:
                # Try with root
                ok, _ = run_cmd(["bash", "-c", f"echo 1 > {profile_path}"])
                if ok:
                    result.ok("amd_power_profile=3D_FULL_SCREEN (via shell)")
                else:
                    result.fail("amd_power_profile", "cannot write profile mode")

    def _apply_amd_performance_level(self, result: OptimizationResult, dry_run: bool):
        """Set AMD GPU to manual or high performance level.

        power_dpm_force_performance_level controls the GPU's power state
        management:
        - auto: driver manages (default)
        - high: lock to highest performance state
        - manual: allow manual clock control
        - profile_peak: maximum clocks
        """
        if dry_run:
            result.skip("amd_perf_level", "dry_run")
            return

        device_paths = self._get_amd_device_paths()
        for dev in device_paths:
            level_path = f"{dev}/power_dpm_force_performance_level"
            if not Path(level_path).exists():
                continue

            orig = sysfs_read(level_path)
            self._rollback.save_original("amd_perf_level", orig)

            if orig == "high" or orig == "profile_peak":
                result.skip("amd_perf_level", f"already {orig}")
                return

            # Try profile_peak first (highest clocks), fall back to high
            if sysfs_write(level_path, "profile_peak"):
                result.ok("amd_perf_level=profile_peak")
            elif sysfs_write(level_path, "high"):
                result.ok("amd_perf_level=high")
            else:
                ok, _ = run_cmd(["bash", "-c", f"echo high > {level_path}"])
                if ok:
                    result.ok("amd_perf_level=high (via shell)")
                else:
                    result.fail("amd_perf_level", "cannot set performance level")

    def _apply_amd_fan_control(self, result: OptimizationResult, dry_run: bool):
        """Set AMD GPU fan to performance curve.

        Keeping thermals low prevents thermal throttling, which causes
        frame drops and latency spikes.
        """
        if dry_run:
            result.skip("amd_fan", "dry_run")
            return

        hwmon_paths = glob.glob(AMDGPU_HWMON)
        for hwmon in hwmon_paths:
            enable_path = f"{hwmon}/pwm1_enable"
            if not Path(enable_path).exists():
                continue

            orig = sysfs_read(enable_path)
            self._rollback.save_original("amd_fan_enable", orig)

            # 1 = manual, 2 = auto (default)
            # We keep auto but ensure it's enabled (some configs disable it)
            if orig != "2":
                sysfs_write(enable_path, "2")

            result.skip("amd_fan", "keeping auto fan control (safe default)")
            return

    def _apply_amd_clock_boost(self, result: OptimizationResult, dry_run: bool):
        """Verify AMD GPU is running at maximum clocks."""
        if dry_run:
            result.skip("amd_clocks", "dry_run")
            return

        device_paths = self._get_amd_device_paths()
        for dev in device_paths:
            sclk_path = f"{dev}/pp_dpm_sclk"
            if not Path(sclk_path).exists():
                continue

            clocks = sysfs_read(sclk_path)
            if clocks:
                # The highest clock level is the last line; '*' marks active
                lines = clocks.strip().splitlines()
                if lines:
                    last_level = lines[-1].split(":")[0].strip()
                    result.ok(f"amd_clocks=verified ({len(lines)} levels, checking highest)")
            return

        result.skip("amd_clocks", "no clock control found")

    def _restore_amd(self, result: OptimizationResult):
        restored = 0

        # Restore performance level
        orig = self._rollback.get_original("amd_perf_level")
        if orig:
            device_paths = self._get_amd_device_paths()
            for dev in device_paths:
                level_path = f"{dev}/power_dpm_force_performance_level"
                if sysfs_write(level_path, orig):
                    restored += 1

        # Restore power profile
        orig = self._rollback.get_original("amd_power_profile")
        if orig:
            device_paths = self._get_amd_device_paths()
            for dev in device_paths:
                profile_path = f"{dev}/pp_power_profile_mode"
                # Profile mode output includes the table; to restore default write "0"
                if sysfs_write(profile_path, "0"):
                    restored += 1

        if restored:
            result.ok(f"amd gpu restored ({restored} settings)")

    # ── NVIDIA GPU (nvidia-smi) ──────────────────────────────────────────

    def _apply_nvidia_perf_mode(self, result: OptimizationResult, dry_run: bool):
        """Set NVIDIA GPU to maximum performance mode.

        nvidia-smi -pm 1: enable persistent mode (keeps driver loaded)
        nvidia-smi -ac: set application clocks to max
        """
        if dry_run:
            result.skip("nvidia_perf", "dry_run")
            return

        # Enable persistence mode
        ok, _ = run_cmd(["nvidia-smi", "-pm", "1"])
        if ok:
            self._rollback.save_original("nvidia_pm", "0")

        # Query max clocks
        ok, output = run_cmd([
            "nvidia-smi",
            "--query-gpu=clocks.max.sm,clocks.max.mem",
            "--format=csv,noheader,nounits"
        ])
        if ok and output:
            parts = output.strip().split(",")
            if len(parts) == 2:
                sm_max = parts[0].strip()
                mem_max = parts[1].strip()

                # Save current clocks
                ok2, current = run_cmd([
                    "nvidia-smi",
                    "--query-gpu=clocks.sm,clocks.mem",
                    "--format=csv,noheader,nounits"
                ])
                if ok2:
                    self._rollback.save_original("nvidia_clocks", current.strip())

                # Set application clocks to max
                ok3, _ = run_cmd([
                    "nvidia-smi", "-ac", f"{mem_max},{sm_max}"
                ])
                if ok3:
                    result.ok(f"nvidia_perf=max clocks (SM:{sm_max}MHz, Mem:{mem_max}MHz)")
                    return

        # Fallback: at least set compute mode to exclusive
        ok, _ = run_cmd(["nvidia-smi", "-c", "EXCLUSIVE_PROCESS"])
        if ok:
            result.ok("nvidia_perf=exclusive_process mode")
        else:
            result.skip("nvidia_perf", "nvidia-smi not available or no changes")

    def _apply_nvidia_power_limit(self, result: OptimizationResult, dry_run: bool):
        """Set NVIDIA GPU power limit to maximum."""
        if not is_admin():
            result.skip("nvidia_power_limit", "not root")
            return
        if dry_run:
            result.skip("nvidia_power_limit", "dry_run")
            return

        # Query max power limit
        ok, output = run_cmd([
            "nvidia-smi",
            "--query-gpu=power.max_limit",
            "--format=csv,noheader,nounits"
        ])
        if not ok:
            result.skip("nvidia_power_limit", "nvidia-smi not available")
            return

        max_power = output.strip()
        if not max_power:
            return

        # Save current power limit
        ok, current = run_cmd([
            "nvidia-smi",
            "--query-gpu=power.limit",
            "--format=csv,noheader,nounits"
        ])
        if ok:
            self._rollback.save_original("nvidia_power_limit", current.strip())

        ok, _ = run_cmd(["nvidia-smi", "-pl", max_power])
        if ok:
            result.ok(f"nvidia_power_limit={max_power}W (maximum)")
        else:
            result.fail("nvidia_power_limit", "could not set power limit")

    def _restore_nvidia(self, result: OptimizationResult):
        restored = 0

        # Restore persistence mode
        orig_pm = self._rollback.get_original("nvidia_pm")
        if orig_pm:
            run_cmd(["nvidia-smi", "-pm", orig_pm])
            restored += 1

        # Restore clocks
        run_cmd(["nvidia-smi", "-rac"])  # Reset application clocks
        restored += 1

        # Restore power limit
        orig_power = self._rollback.get_original("nvidia_power_limit")
        if orig_power:
            run_cmd(["nvidia-smi", "-pl", orig_power])
            restored += 1

        # Restore compute mode
        run_cmd(["nvidia-smi", "-c", "DEFAULT"])

        if restored:
            result.ok(f"nvidia gpu restored ({restored} settings)")

    # ── Status ───────────────────────────────────────────────────────────

    def get_status(self) -> dict:
        status = {"module": "linux_gpu", "optimized": False, "settings": {}}

        if not IS_LINUX:
            return status

        vendor = self._detect_gpu_vendor()
        status["settings"]["vendor"] = vendor

        if "amd" in (vendor or ""):
            device_paths = self._get_amd_device_paths()
            for dev in device_paths:
                # Performance level
                level = sysfs_read(f"{dev}/power_dpm_force_performance_level")
                if level:
                    status["settings"]["amd_perf_level"] = level

                # Current SCLK
                sclk = sysfs_read(f"{dev}/pp_dpm_sclk")
                if sclk:
                    for line in sclk.splitlines():
                        if "*" in line:
                            status["settings"]["amd_current_sclk"] = line.strip()

                # GPU temp
                hwmon_temps = glob.glob(f"{dev}/hwmon/hwmon*/temp1_input")
                for t in hwmon_temps:
                    temp = sysfs_read(t)
                    if temp:
                        status["settings"]["gpu_temp_c"] = int(temp) // 1000

                # Power usage
                hwmon_power = glob.glob(f"{dev}/hwmon/hwmon*/power1_average")
                for p in hwmon_power:
                    power = sysfs_read(p)
                    if power:
                        status["settings"]["gpu_power_w"] = int(power) / 1_000_000

        if "nvidia" in (vendor or ""):
            ok, output = run_cmd([
                "nvidia-smi",
                "--query-gpu=name,temperature.gpu,power.draw,clocks.sm,clocks.mem,utilization.gpu",
                "--format=csv,noheader,nounits"
            ])
            if ok and output:
                # nvidia-smi may return multiple lines for multiple GPUs
                for line in output.strip().splitlines():
                    parts = [p.strip() for p in line.split(",")]
                    if len(parts) < 6:
                        continue
                    entry = {"nvidia_name": parts[0]}
                    for key, idx, conv in [
                        ("gpu_temp_c", 1, int), ("gpu_power_w", 2, float),
                        ("nvidia_sm_clock", 3, int), ("nvidia_mem_clock", 4, int),
                        ("gpu_utilization", 5, int),
                    ]:
                        try:
                            if parts[idx] not in ("[N/A]", "N/A", ""):
                                entry[key] = conv(parts[idx])
                        except (ValueError, IndexError):
                            pass
                    gpus = status["settings"].setdefault("nvidia_gpus", [])
                    gpus.append(entry)
                    # Also set top-level from first GPU
                    if "nvidia_name" not in status["settings"]:
                        status["settings"].update(entry)

        status["optimized"] = self._rollback.has_saved_state
        return status

    # ── Helpers ───────────────────────────────────────────────────────────

    def _detect_gpu_vendor(self) -> Optional[str]:
        """Detect GPU vendor(s) from /sys or lspci."""
        vendors = []

        # Check for AMD GPU via sysfs
        amd_paths = glob.glob("/sys/class/drm/card*/device/vendor")
        for vp in amd_paths:
            vendor_id = sysfs_read(vp)
            if vendor_id == "0x1002":  # AMD
                vendors.append("amd")
            elif vendor_id == "0x10de":  # NVIDIA
                vendors.append("nvidia")

        if vendors:
            return "+".join(sorted(set(vendors)))

        # Fallback to lspci
        ok, output = run_cmd(["lspci", "-nn"])
        if ok:
            lower = output.lower()
            if "amd" in lower or "radeon" in lower or "[1002:" in lower:
                vendors.append("amd")
            if "nvidia" in lower or "geforce" in lower or "[10de:" in lower:
                vendors.append("nvidia")

        return "+".join(sorted(set(vendors))) if vendors else None

    def _get_amd_device_paths(self) -> list[str]:
        """Get paths to AMD GPU device sysfs directories."""
        paths = []
        for vp in glob.glob("/sys/class/drm/card*/device/vendor"):
            vendor_id = sysfs_read(vp)
            if vendor_id == "0x1002":
                paths.append(str(Path(vp).parent))
        return paths
