"""Linux network latency optimization via sysctl and ethtool."""
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__)


class LinuxNetworkOptimizer:
    """Optimize Linux network stack for minimum latency."""

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

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

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

        self._apply_tcp_tuning(result, dry_run)
        self._apply_tcp_fastopen(result, dry_run)
        self._apply_bbr(result, dry_run)
        self._apply_adapter_tuning(result, dry_run)
        self._apply_irq_balance(result, dry_run)

        return result

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

        self._restore_tcp_tuning(result)
        self._restore_tcp_fastopen(result)
        self._restore_bbr(result)
        self._restore_adapter_tuning(result)

        self._rollback.clear()
        return result

    # ── TCP Stack Tuning ─────────────────────────────────────────────────

    def _apply_tcp_tuning(self, result: OptimizationResult, dry_run: bool):
        """Tune TCP stack parameters for low latency.

        - tcp_nodelay via application hints (sysctl can't force this)
        - tcp_low_latency: hint to prefer latency over throughput
        - tcp_timestamps: disable to save 12 bytes/packet
        - tcp_sack: disable for lower overhead (acceptable on good networks)
        - Larger socket buffers for high-throughput scenarios
        """
        if not is_admin():
            result.skip("tcp_tuning", "not root")
            return
        if dry_run:
            result.skip("tcp_tuning", "dry_run")
            return

        params = {
            # Prefer low latency over throughput
            "/proc/sys/net/ipv4/tcp_low_latency": "1",
            # Disable timestamps (12 bytes/packet savings)
            "/proc/sys/net/ipv4/tcp_timestamps": "0",
            # Disable SACK (lower overhead, fine on reliable networks)
            "/proc/sys/net/ipv4/tcp_sack": "0",
            # Shorter TIME_WAIT recycling
            "/proc/sys/net/ipv4/tcp_tw_reuse": "1",
            # Shorter FIN timeout
            "/proc/sys/net/ipv4/tcp_fin_timeout": "15",
            # Disable slow start after idle
            "/proc/sys/net/ipv4/tcp_slow_start_after_idle": "0",
            # Increase max backlog for high packet rates
            "/proc/sys/net/core/netdev_max_backlog": "5000",
            # Larger socket buffers
            "/proc/sys/net/core/rmem_max": "16777216",
            "/proc/sys/net/core/wmem_max": "16777216",
            # TCP buffer autotuning range (min, default, max)
            "/proc/sys/net/ipv4/tcp_rmem": "4096 87380 16777216",
            "/proc/sys/net/ipv4/tcp_wmem": "4096 65536 16777216",
            # Increase somaxconn for better connection handling
            "/proc/sys/net/core/somaxconn": "4096",
        }

        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"tcp_{key}", orig)
                if sysfs_write(path, value):
                    changed += 1

        if changed:
            result.ok(f"tcp_tuning={changed} params (low latency, no timestamps, larger buffers)")
        else:
            result.fail("tcp_tuning", "no params changed")

    def _restore_tcp_tuning(self, result: OptimizationResult):
        params = [
            "tcp_low_latency", "tcp_timestamps", "tcp_sack",
            "tcp_tw_reuse", "tcp_fin_timeout", "tcp_slow_start_after_idle",
            "netdev_max_backlog", "rmem_max", "wmem_max",
            "tcp_rmem", "tcp_wmem", "somaxconn",
        ]
        restored = 0
        for name in params:
            orig = self._rollback.get_original(f"tcp_{name}")
            if orig is not None:
                if name in ("rmem_max", "wmem_max", "netdev_max_backlog", "somaxconn"):
                    path = f"/proc/sys/net/core/{name}"
                else:
                    path = f"/proc/sys/net/ipv4/{name}"
                if sysfs_write(path, orig):
                    restored += 1
        if restored:
            result.ok(f"tcp_tuning restored ({restored} params)")

    # ── TCP Fast Open ────────────────────────────────────────────────────

    def _apply_tcp_fastopen(self, result: OptimizationResult, dry_run: bool):
        """Enable TCP Fast Open for client and server.

        TFO saves one RTT on connection establishment by sending data
        in the SYN packet. Value 3 = enabled for both client and server.
        """
        if not is_admin():
            result.skip("tcp_fastopen", "not root")
            return
        if dry_run:
            result.skip("tcp_fastopen", "dry_run")
            return

        path = "/proc/sys/net/ipv4/tcp_fastopen"
        if not Path(path).exists():
            result.skip("tcp_fastopen", "not available")
            return

        orig = sysfs_read(path)
        self._rollback.save_original("tcp_fastopen", orig)

        if orig == "3":
            result.skip("tcp_fastopen", "already enabled (client+server)")
            return

        if sysfs_write(path, "3"):
            result.ok(f"tcp_fastopen=3 (client+server, was {orig})")
        else:
            result.fail("tcp_fastopen")

    def _restore_tcp_fastopen(self, result: OptimizationResult):
        orig = self._rollback.get_original("tcp_fastopen")
        if orig is not None:
            if sysfs_write("/proc/sys/net/ipv4/tcp_fastopen", orig):
                result.ok("tcp_fastopen restored")

    # ── BBR Congestion Control ───────────────────────────────────────────

    def _apply_bbr(self, result: OptimizationResult, dry_run: bool):
        """Switch to BBR congestion control if available.

        BBR (Bottleneck Bandwidth and RTT) is Google's congestion
        control algorithm. It's better than CUBIC for:
        - Lower latency under congestion
        - Better throughput on lossy networks
        - Faster recovery from packet loss
        """
        if not is_admin():
            result.skip("bbr", "not root")
            return
        if dry_run:
            result.skip("bbr", "dry_run")
            return

        # Check if BBR module is available
        avail_path = "/proc/sys/net/ipv4/tcp_available_congestion_control"
        available = sysfs_read(avail_path)
        if not available or "bbr" not in available:
            # Try to load the module
            run_cmd(["modprobe", "tcp_bbr"])
            available = sysfs_read(avail_path)
            if not available or "bbr" not in available:
                result.skip("bbr", "BBR module not available")
                return

        cc_path = "/proc/sys/net/ipv4/tcp_congestion_control"
        orig = sysfs_read(cc_path)
        self._rollback.save_original("tcp_cc", orig)

        if orig == "bbr":
            result.skip("bbr", "already using BBR")
            return

        if sysfs_write(cc_path, "bbr"):
            # Also set fq as the qdisc (BBR works best with fq)
            run_cmd(["tc", "qdisc", "replace", "dev", "eth0", "root", "fq"])
            result.ok(f"bbr=enabled (was {orig})")
        else:
            result.fail("bbr", "could not set congestion control")

    def _restore_bbr(self, result: OptimizationResult):
        orig = self._rollback.get_original("tcp_cc")
        if orig is not None:
            if sysfs_write("/proc/sys/net/ipv4/tcp_congestion_control", orig):
                result.ok(f"congestion control restored to {orig}")

    # ── NIC Adapter Tuning ───────────────────────────────────────────────

    def _apply_adapter_tuning(self, result: OptimizationResult, dry_run: bool):
        """Tune NIC settings via ethtool for low latency.

        - Disable interrupt coalescing (lower latency, higher CPU)
        - Increase ring buffer sizes
        - Enable GRO (Generic Receive Offload) for throughput
        """
        if not is_admin():
            result.skip("adapter_tuning", "not root")
            return
        if dry_run:
            result.skip("adapter_tuning", "dry_run")
            return

        interfaces = self._get_active_interfaces()
        if not interfaces:
            result.skip("adapter_tuning", "no active interfaces found")
            return

        changed = 0
        for iface in interfaces:
            # Disable interrupt coalescing for lowest latency
            ok, _ = run_cmd([
                "ethtool", "-C", iface,
                "rx-usecs", "0", "tx-usecs", "0",
                "rx-frames", "1", "tx-frames", "1",
            ])
            if ok:
                changed += 1
                self._rollback.save_original(f"adapter_{iface}_coalesce", "default")

            # Maximize ring buffers
            ok, output = run_cmd(["ethtool", "-g", iface])
            if ok:
                # Parse max ring buffer sizes and set them
                import re
                rx_max = re.search(r"RX:\s+(\d+)", output)
                tx_max = re.search(r"TX:\s+(\d+)", output)
                if rx_max and tx_max:
                    run_cmd([
                        "ethtool", "-G", iface,
                        "rx", rx_max.group(1),
                        "tx", tx_max.group(1),
                    ])
                    changed += 1

            # Enable GRO and TSO
            run_cmd(["ethtool", "-K", iface, "gro", "on", "tso", "on", "gso", "on"])

        if changed:
            result.ok(f"adapter_tuning={changed} changes on {len(interfaces)} interfaces")
        else:
            result.skip("adapter_tuning", "ethtool not available or no changes")

    def _restore_adapter_tuning(self, result: OptimizationResult):
        interfaces = self._get_active_interfaces()
        restored = 0
        for iface in interfaces:
            if self._rollback.has_original(f"adapter_{iface}_coalesce"):
                # Restore default coalescing (adaptive)
                ok, _ = run_cmd(["ethtool", "-C", iface, "adaptive-rx", "on", "adaptive-tx", "on"])
                if ok:
                    restored += 1
        if restored:
            result.ok(f"adapter_tuning restored ({restored} interfaces)")

    # ── IRQ Balancing ────────────────────────────────────────────────────

    def _apply_irq_balance(self, result: OptimizationResult, dry_run: bool):
        """Stop irqbalance and manually pin NIC IRQs.

        irqbalance is good for general use but adds latency variance.
        For gaming/low-latency, stopping it and pinning IRQs to specific
        cores gives more consistent performance.
        """
        if not is_admin():
            result.skip("irq_balance", "not root")
            return
        if dry_run:
            result.skip("irq_balance", "dry_run")
            return

        # Check if irqbalance is running
        ok, output = run_cmd(["systemctl", "is-active", "irqbalance"])
        if ok and "active" in output:
            self._rollback.save_original("irqbalance_was_active", True)
            ok, _ = run_cmd(["systemctl", "stop", "irqbalance"])
            if ok:
                result.ok("irq_balance=stopped irqbalance for consistent latency")
            else:
                result.fail("irq_balance", "could not stop irqbalance")
        else:
            result.skip("irq_balance", "irqbalance not running")

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

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

        if not IS_LINUX:
            return status

        # TCP congestion control
        cc = sysfs_read("/proc/sys/net/ipv4/tcp_congestion_control")
        if cc:
            status["settings"]["congestion_control"] = cc

        # TCP Fast Open
        tfo = sysfs_read("/proc/sys/net/ipv4/tcp_fastopen")
        if tfo:
            status["settings"]["tcp_fastopen"] = int(tfo)

        # Low latency mode
        ll = sysfs_read("/proc/sys/net/ipv4/tcp_low_latency")
        if ll:
            status["settings"]["tcp_low_latency"] = bool(int(ll))

        # Timestamps
        ts = sysfs_read("/proc/sys/net/ipv4/tcp_timestamps")
        if ts:
            status["settings"]["tcp_timestamps"] = bool(int(ts))

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

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

    def _get_active_interfaces(self) -> list[str]:
        """Get list of active (UP) non-loopback network interfaces."""
        interfaces = []
        net_dir = Path("/sys/class/net")
        if not net_dir.exists():
            return interfaces

        for iface_dir in net_dir.iterdir():
            name = iface_dir.name
            if name == "lo":
                continue
            operstate = sysfs_read(str(iface_dir / "operstate"))
            if operstate == "up":
                interfaces.append(name)

        return interfaces
