"""Linux power and CPU scheduler optimization."""
from __future__ import annotations

import glob
import logging
import os
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__)

# CPU frequency governor sysfs paths
CPUFREQ_BASE = "/sys/devices/system/cpu"
INTEL_PSTATE = "/sys/devices/system/cpu/intel_pstate"
AMD_PSTATE = "/sys/devices/system/cpu/amd_pstate"


class LinuxPowerOptimizer:
    """Optimize Linux power and CPU settings for maximum performance."""

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

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

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

        self._apply_cpu_governor(result, dry_run)
        self._apply_cpu_boost(result, dry_run)
        self._apply_energy_perf(result, dry_run)
        self._apply_scheduler(result, dry_run)
        self._apply_kernel_latency(result, dry_run)

        return result

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

        self._restore_cpu_governor(result)
        self._restore_cpu_boost(result)
        self._restore_energy_perf(result)
        self._restore_scheduler(result)
        self._restore_kernel_latency(result)

        self._rollback.clear()
        return result

    # ── CPU Governor ─────────────────────────────────────────────────────

    def _apply_cpu_governor(self, result: OptimizationResult, dry_run: bool):
        """Set CPU frequency governor to 'performance' on all cores.

        The 'performance' governor locks CPUs at max frequency — no
        frequency scaling overhead, no ramp-up latency. Uses more power
        but eliminates the 1-5ms latency from frequency transitions.
        """
        if not is_admin():
            result.skip("cpu_governor", "not root")
            return
        if dry_run:
            result.skip("cpu_governor", "dry_run")
            return

        cpu_dirs = sorted(glob.glob(f"{CPUFREQ_BASE}/cpu[0-9]*/cpufreq/scaling_governor"))
        if not cpu_dirs:
            result.skip("cpu_governor", "no cpufreq sysfs found")
            return

        # Save first core's governor as reference
        orig = sysfs_read(cpu_dirs[0])
        self._rollback.save_original("cpu_governor", orig)

        if orig == "performance":
            result.skip("cpu_governor", "already performance")
            return

        changed = 0
        for gov_path in cpu_dirs:
            if sysfs_write(gov_path, "performance"):
                changed += 1

        if changed:
            result.ok(f"cpu_governor=performance on {changed} cores (was {orig})")
        else:
            result.fail("cpu_governor", "failed to set governor")

    def _restore_cpu_governor(self, result: OptimizationResult):
        orig = self._rollback.get_original("cpu_governor")
        if not orig:
            return
        cpu_dirs = sorted(glob.glob(f"{CPUFREQ_BASE}/cpu[0-9]*/cpufreq/scaling_governor"))
        restored = 0
        for gov_path in cpu_dirs:
            if sysfs_write(gov_path, orig):
                restored += 1
        if restored:
            result.ok(f"cpu_governor restored to {orig} on {restored} cores")

    # ── CPU Boost (Turbo) ────────────────────────────────────────────────

    def _apply_cpu_boost(self, result: OptimizationResult, dry_run: bool):
        """Enable CPU turbo boost.

        On Intel: /sys/devices/system/cpu/intel_pstate/no_turbo = 0
        On AMD: /sys/devices/system/cpu/cpufreq/boost = 1
        Generic: /sys/devices/system/cpu/cpufreq/boost = 1
        """
        if not is_admin():
            result.skip("cpu_boost", "not root")
            return
        if dry_run:
            result.skip("cpu_boost", "dry_run")
            return

        # Intel pstate
        intel_turbo = f"{INTEL_PSTATE}/no_turbo"
        if Path(intel_turbo).exists():
            orig = sysfs_read(intel_turbo)
            self._rollback.save_original("intel_no_turbo", orig)
            if orig == "0":
                result.skip("cpu_boost", "turbo already enabled")
                return
            if sysfs_write(intel_turbo, "0"):
                result.ok("cpu_boost=turbo enabled (Intel)")
                return

        # AMD/generic boost
        boost_path = f"{CPUFREQ_BASE}/cpufreq/boost"
        if Path(boost_path).exists():
            orig = sysfs_read(boost_path)
            self._rollback.save_original("cpu_boost", orig)
            if orig == "1":
                result.skip("cpu_boost", "boost already enabled")
                return
            if sysfs_write(boost_path, "1"):
                result.ok("cpu_boost=enabled")
                return

        result.skip("cpu_boost", "no boost control found")

    def _restore_cpu_boost(self, result: OptimizationResult):
        intel_orig = self._rollback.get_original("intel_no_turbo")
        if intel_orig is not None:
            intel_turbo = f"{INTEL_PSTATE}/no_turbo"
            if sysfs_write(intel_turbo, intel_orig):
                result.ok("cpu_boost restored (Intel)")
                return

        boost_orig = self._rollback.get_original("cpu_boost")
        if boost_orig is not None:
            boost_path = f"{CPUFREQ_BASE}/cpufreq/boost"
            if sysfs_write(boost_path, boost_orig):
                result.ok("cpu_boost restored")

    # ── Energy Performance Preference ────────────────────────────────────

    def _apply_energy_perf(self, result: OptimizationResult, dry_run: bool):
        """Set energy performance preference to 'performance' on all cores.

        This controls the hardware's power/performance tradeoff at the
        firmware level (EPP on Intel, EPP on AMD with amd_pstate).
        """
        if not is_admin():
            result.skip("energy_perf", "not root")
            return
        if dry_run:
            result.skip("energy_perf", "dry_run")
            return

        epp_paths = sorted(glob.glob(
            f"{CPUFREQ_BASE}/cpu[0-9]*/cpufreq/energy_performance_preference"
        ))
        if not epp_paths:
            result.skip("energy_perf", "EPP not available")
            return

        orig = sysfs_read(epp_paths[0])
        self._rollback.save_original("energy_perf_pref", orig)

        if orig == "performance":
            result.skip("energy_perf", "already performance")
            return

        changed = 0
        for epp in epp_paths:
            if sysfs_write(epp, "performance"):
                changed += 1

        if changed:
            result.ok(f"energy_perf=performance on {changed} cores (was {orig})")
        else:
            result.fail("energy_perf", "could not set EPP")

    def _restore_energy_perf(self, result: OptimizationResult):
        orig = self._rollback.get_original("energy_perf_pref")
        if not orig:
            return
        epp_paths = sorted(glob.glob(
            f"{CPUFREQ_BASE}/cpu[0-9]*/cpufreq/energy_performance_preference"
        ))
        restored = 0
        for epp in epp_paths:
            if sysfs_write(epp, orig):
                restored += 1
        if restored:
            result.ok(f"energy_perf restored to {orig}")

    # ── Kernel Scheduler Tuning ──────────────────────────────────────────

    def _apply_scheduler(self, result: OptimizationResult, dry_run: bool):
        """Tune kernel scheduler for low latency.

        - sched_min_granularity_ns: minimum timeslice (lower = more responsive)
        - sched_wakeup_granularity_ns: wakeup preemption threshold
        - sched_migration_cost_ns: how quickly threads migrate between cores
        """
        if not is_admin():
            result.skip("scheduler", "not root")
            return
        if dry_run:
            result.skip("scheduler", "dry_run")
            return

        sched_params = {
            # Shorter timeslices = more responsive to new work
            "/proc/sys/kernel/sched_min_granularity_ns": "500000",      # 0.5ms (default ~3ms)
            "/proc/sys/kernel/sched_wakeup_granularity_ns": "500000",   # 0.5ms
            # Lower migration cost = faster thread rebalancing
            "/proc/sys/kernel/sched_migration_cost_ns": "250000",       # 0.25ms (default 0.5ms)
            # Disable autogroup (interferes with manual nice/priority)
            "/proc/sys/kernel/sched_autogroup_enabled": "0",
        }

        changed = 0
        for path, value in sched_params.items():
            if Path(path).exists():
                orig = sysfs_read(path)
                key = path.split("/")[-1]
                self._rollback.save_original(f"sched_{key}", orig)
                if sysfs_write(path, value):
                    changed += 1

        if changed:
            result.ok(f"scheduler={changed} params tuned (low-latency timeslices)")
        else:
            result.skip("scheduler", "no params changed")

    def _restore_scheduler(self, result: OptimizationResult):
        sched_files = [
            "sched_min_granularity_ns",
            "sched_wakeup_granularity_ns",
            "sched_migration_cost_ns",
            "sched_autogroup_enabled",
        ]
        restored = 0
        for name in sched_files:
            orig = self._rollback.get_original(f"sched_{name}")
            if orig is not None:
                path = f"/proc/sys/kernel/{name}"
                if sysfs_write(path, orig):
                    restored += 1
        if restored:
            result.ok(f"scheduler restored ({restored} params)")

    # ── Kernel Latency Tuning ────────────────────────────────────────────

    def _apply_kernel_latency(self, result: OptimizationResult, dry_run: bool):
        """Tune kernel parameters for minimum latency.

        - vm.swappiness: reduce swap pressure
        - vm.dirty_ratio: flush writes sooner
        - kernel.timer_migration: disable (keep timers on same core)
        """
        if not is_admin():
            result.skip("kernel_latency", "not root")
            return
        if dry_run:
            result.skip("kernel_latency", "dry_run")
            return

        params = {
            "/proc/sys/vm/swappiness": "10",              # Default 60, lower = less swap
            "/proc/sys/vm/dirty_ratio": "10",              # Flush at 10% dirty (default 20)
            "/proc/sys/vm/dirty_background_ratio": "5",    # Start background flush at 5%
            "/proc/sys/vm/dirty_expire_centisecs": "500",  # Flush dirty pages after 5s
            "/proc/sys/kernel/timer_migration": "0",       # Keep timers on same core
        }

        changed = 0
        for path, value in params.items():
            if Path(path).exists():
                orig = sysfs_read(path)
                key = path.split("/")[-1]
                self._rollback.save_original(f"kern_{key}", orig)
                if sysfs_write(path, value):
                    changed += 1

        if changed:
            result.ok(f"kernel_latency={changed} params (low swappiness, fast flush, no timer migration)")
        else:
            result.skip("kernel_latency", "no params changed")

    def _restore_kernel_latency(self, result: OptimizationResult):
        params = ["swappiness", "dirty_ratio", "dirty_background_ratio",
                  "dirty_expire_centisecs", "timer_migration"]
        restored = 0
        for name in params:
            orig = self._rollback.get_original(f"kern_{name}")
            if orig is not None:
                # Map back to correct /proc path
                if name.startswith("dirty") or name == "swappiness":
                    path = f"/proc/sys/vm/{name}"
                else:
                    path = f"/proc/sys/kernel/{name}"
                if sysfs_write(path, orig):
                    restored += 1
        if restored:
            result.ok(f"kernel_latency restored ({restored} params)")

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

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

        if not IS_LINUX:
            return status

        # Governor
        gov_path = f"{CPUFREQ_BASE}/cpu0/cpufreq/scaling_governor"
        gov = sysfs_read(gov_path)
        if gov:
            status["settings"]["cpu_governor"] = gov

        # Current frequency
        freq_path = f"{CPUFREQ_BASE}/cpu0/cpufreq/scaling_cur_freq"
        freq = sysfs_read(freq_path)
        if freq:
            status["settings"]["cpu_freq_mhz"] = int(freq) // 1000

        # Max frequency
        max_freq_path = f"{CPUFREQ_BASE}/cpu0/cpufreq/scaling_max_freq"
        max_freq = sysfs_read(max_freq_path)
        if max_freq:
            status["settings"]["cpu_max_freq_mhz"] = int(max_freq) // 1000

        # Swappiness
        swap = sysfs_read("/proc/sys/vm/swappiness")
        if swap:
            status["settings"]["swappiness"] = int(swap)

        # Scheduler
        gran = sysfs_read("/proc/sys/kernel/sched_min_granularity_ns")
        if gran:
            status["settings"]["sched_min_granularity_ms"] = int(gran) / 1_000_000

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