# Copyright (c) 2025 Muhammad Nawaz <m.nawaz2003@gmail.com>
# SPDX-License-Identifier: MIT
"""User-space stack configuration generation for Onload, VMA, and RDMA.

This module generates environment profiles and configuration files for
user-space networking stacks commonly used in HFT environments.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class Issue:
    """Validation issue for user stack configuration."""

    severity: str  # info/warning/error
    component: str  # onload/vma/rdma
    message: str


def generate_onload_profile(
    nic_role: str = "trading",
    tuning_level: str = "ultra_low_latency",
    interface: Optional[str] = None,
) -> str:
    """Generate Onload environment profile for shell sourcing.

    Args:
        nic_role: Role of the NIC (trading/control/multicast)
        tuning_level: Tuning level (ultra_low_latency/low_latency/balanced)
        interface: Specific interface to use (optional)

    Returns:
        Shell script content for sourcing
    """
    # Base settings common to all profiles
    base_settings = {
        "EF_POLL_USEC": "0",  # Busy poll (0 = adaptive)
        "EF_INT_DRIVEN": "0",  # Use spinning, not interrupts
        "EF_SPIN_USEC": "-1",  # Spin forever
        "EF_HELPER_USEC": "0",  # Disable helper thread
        "EF_NONAGLE_INFLIGHT_MAX": "1",  # Disable Nagle
        "EF_TCP_SNDBUF": "1048576",  # 1MB send buffer
        "EF_TCP_RCVBUF": "1048576",  # 1MB receive buffer
        "EF_UDP_SNDBUF": "1048576",
        "EF_UDP_RCVBUF": "1048576",
    }

    # Ultra-low-latency specific settings
    if tuning_level == "ultra_low_latency":
        base_settings.update(
            {
                "EF_POLL_USEC": "0",
                "EF_SPIN_USEC": "-1",  # Infinite spin
                "EF_SLEEP_SPIN_USEC": "0",  # No sleep
                "EF_PKT_WAIT_SPIN": "1",  # Spin waiting for packets
                "EF_TCP_FASTSTART": "0",  # Disable slow start
                "EF_HIGH_THROUGHPUT_MODE": "0",  # Latency over throughput
                "EF_IRQ_MODERATION": "0",  # Disable IRQ moderation
                "EF_EVS_PER_POLL": "128",  # Events per poll
            }
        )
    elif tuning_level == "low_latency":
        base_settings.update(
            {
                "EF_POLL_USEC": "100",  # Some adaptive polling
                "EF_SPIN_USEC": "100000",  # 100ms spin then block
                "EF_HIGH_THROUGHPUT_MODE": "0",
            }
        )
    elif tuning_level == "balanced":
        base_settings.update(
            {
                "EF_POLL_USEC": "1000",  # More aggressive sleeping
                "EF_SPIN_USEC": "10000",  # 10ms spin
                "EF_INT_DRIVEN": "1",  # Use interrupts
            }
        )

    # Role-specific adjustments
    if nic_role == "multicast":
        base_settings.update(
            {
                "EF_MCAST_JOIN_HANDOVER": "1",
                "EF_MCAST_RECV": "1",
                "EF_FORCE_SEND_MULTICAST": "1",
            }
        )
    elif nic_role == "control":
        # Control path doesn't need ultra-low latency
        base_settings.update(
            {
                "EF_POLL_USEC": "1000",
                "EF_INT_DRIVEN": "1",
            }
        )

    # Interface binding
    if interface:
        base_settings["EF_INTERFACE"] = interface

    # Generate shell script
    lines = [
        "#!/bin/bash",
        "# Onload environment profile generated by lltune",
        f"# Role: {nic_role}, Tuning: {tuning_level}",
        "#",
        "# Source this file before running your application:",
        "#   source /etc/profile.d/lltune-onload.sh",
        "#   onload ./your_trading_app",
        "",
        "# Onload environment variables",
    ]

    for key, value in sorted(base_settings.items()):
        lines.append(f'export {key}="{value}"')

    lines.extend(
        [
            "",
            "# Verify Onload is available",
            "if ! command -v onload &> /dev/null; then",
            '    echo "Warning: onload command not found" >&2',
            "fi",
            "",
        ]
    )

    return "\n".join(lines)


def generate_vma_profile(
    nic_role: str = "trading",
    tuning_level: str = "ultra_low_latency",
    interface: Optional[str] = None,
) -> str:
    """Generate Mellanox VMA environment profile for shell sourcing.

    Args:
        nic_role: Role of the NIC (trading/control/multicast)
        tuning_level: Tuning level (ultra_low_latency/low_latency/balanced)
        interface: Specific interface to use (optional)

    Returns:
        Shell script content for sourcing
    """
    # Base VMA settings
    base_settings = {
        "VMA_SPEC": "latency",  # Use latency profile
        "VMA_RING_ALLOCATION_LOGIC_RX": "1",  # Per-socket rings
        "VMA_RING_ALLOCATION_LOGIC_TX": "1",
        "VMA_RX_POLL": "-1",  # Infinite polling
        "VMA_SELECT_POLL": "-1",
        "VMA_THREAD_MODE": "1",  # Multi-threaded
        "VMA_MTU": "1500",  # Standard MTU
        "VMA_RX_BUFS": "200000",  # Large buffer pool
        "VMA_TX_BUFS": "200000",
        "VMA_RX_WRE": "16000",  # Work request elements
        "VMA_TX_WRE": "3000",
    }

    # Tuning level adjustments
    if tuning_level == "ultra_low_latency":
        base_settings.update(
            {
                "VMA_RX_POLL": "-1",  # Infinite spin
                "VMA_SELECT_POLL": "-1",
                "VMA_RING_ALLOCATION_LOGIC_RX": "2",  # Ring per QP
                "VMA_CQ_MODERATION_ENABLE": "0",  # Disable moderation
                "VMA_CQ_AIM_MAX_COUNT": "0",
                "VMA_CQ_AIM_MAX_PERIOD_USEC": "0",
                "VMA_GRO_STREAMS_MAX": "0",  # Disable GRO
                "VMA_BF": "1",  # Enable Blue Flame
            }
        )
    elif tuning_level == "low_latency":
        base_settings.update(
            {
                "VMA_RX_POLL": "100000",  # 100ms polling
                "VMA_SELECT_POLL": "100000",
            }
        )
    elif tuning_level == "balanced":
        base_settings.update(
            {
                "VMA_RX_POLL": "0",  # Interrupt-driven
                "VMA_SELECT_POLL": "0",
                "VMA_CQ_MODERATION_ENABLE": "1",
            }
        )

    # Role-specific adjustments
    if nic_role == "multicast":
        base_settings.update(
            {
                "VMA_IGMP": "2",  # IGMPv2
                "VMA_NEIGH_UC_ARP_DELAY_MSEC": "10000",
            }
        )

    # Interface binding
    if interface:
        base_settings["VMA_INTF"] = interface

    # Generate shell script
    lines = [
        "#!/bin/bash",
        "# VMA (Mellanox Messaging Accelerator) profile generated by lltune",
        f"# Role: {nic_role}, Tuning: {tuning_level}",
        "#",
        "# Source this file before running your application:",
        "#   source /etc/profile.d/lltune-vma.sh",
        "#   LD_PRELOAD=libvma.so ./your_trading_app",
        "",
        "# VMA environment variables",
    ]

    for key, value in sorted(base_settings.items()):
        lines.append(f'export {key}="{value}"')

    lines.extend(
        [
            "",
            "# Convenience function to run with VMA",
            "run_with_vma() {",
            '    LD_PRELOAD="${VMA_LIB:-/usr/lib64/libvma.so}" "$@"',
            "}",
            "",
            "# Check VMA availability",
            'VMA_LIB="${VMA_LIB:-/usr/lib64/libvma.so}"',
            'if [[ ! -f "$VMA_LIB" ]]; then',
            '    echo "Warning: VMA library not found at $VMA_LIB" >&2',
            "fi",
            "",
        ]
    )

    return "\n".join(lines)


def validate_rdma_alignment(
    rdma_devices: List[Dict],
    nics: List[Dict],
    cfg: Optional[Dict] = None,
) -> List[Issue]:
    """Validate RDMA device and NIC NUMA alignment.

    Args:
        rdma_devices: List of RDMA device info dicts
        nics: List of NIC info dicts
        cfg: Optional configuration dict

    Returns:
        List of validation issues found
    """
    issues: List[Issue] = []

    if not rdma_devices:
        return issues

    # Build NUMA map for NICs
    nic_numa: Dict[str, Optional[int]] = {}
    for nic in nics:
        name = nic.get("name")
        numa = nic.get("numa_node")
        if name:
            nic_numa[name] = numa

    # Check each RDMA device
    for rdma in rdma_devices:
        dev_name = rdma.get("name") or rdma.get("device", "unknown")
        rdma_numa = rdma.get("numa_node")

        # Get associated network interfaces
        netdevs = rdma.get("netdevs", [])
        if not netdevs:
            # Try to find by naming convention (mlx5_0 -> related NICs)
            for nic_name in nic_numa:
                if "mlx" in nic_name or "ib" in nic_name:
                    netdevs.append(nic_name)

        for netdev in netdevs:
            nic_node = nic_numa.get(netdev)
            if nic_node is not None and rdma_numa is not None:
                if nic_node != rdma_numa:
                    issues.append(
                        Issue(
                            severity="warning",
                            component="rdma",
                            message=f"RDMA device {dev_name} (NUMA {rdma_numa}) "
                            f"associated with NIC {netdev} (NUMA {nic_node}) - cross-NUMA",
                        )
                    )

    # Check for RDMA devices without proper NUMA placement
    for rdma in rdma_devices:
        dev_name = rdma.get("name") or rdma.get("device", "unknown")
        if rdma.get("numa_node") is None:
            issues.append(
                Issue(
                    severity="info",
                    component="rdma",
                    message=f"RDMA device {dev_name} has no NUMA node information",
                )
            )

    return issues


def generate_rdma_tuning_script(
    rdma_devices: List[Dict],
    cfg: Optional[Dict] = None,
) -> str:
    """Generate RDMA tuning script for mlx5 devices.

    Args:
        rdma_devices: List of RDMA device info
        cfg: Optional configuration

    Returns:
        Shell script for RDMA tuning
    """
    lines = [
        "#!/bin/bash",
        "# RDMA device tuning script generated by lltune",
        "set -euo pipefail",
        "",
        "# This script applies recommended RDMA settings for low-latency workloads",
        "",
    ]

    # mlx5-specific tuning
    lines.extend(
        [
            "# Check for mlx5 module",
            "if ! lsmod | grep -q mlx5_core; then",
            '    echo "mlx5_core module not loaded; skipping mlx5 tuning"',
            "    exit 0",
            "fi",
            "",
            "# Set mlx5 module parameters for low latency",
            "# Note: These require module reload to take effect",
            "MLX5_PARAMS=/sys/module/mlx5_core/parameters",
            "",
            "# Increase completion queue interrupt moderation for polling mode",
            'if [[ -w "$MLX5_PARAMS/comp_eq_depth" ]]; then',
            '    echo 64 > "$MLX5_PARAMS/comp_eq_depth"',
            "fi",
            "",
        ]
    )

    # Device-specific settings
    for rdma in rdma_devices:
        dev_name = rdma.get("name") or rdma.get("device")
        if not dev_name:
            continue

        lines.extend(
            [
                f"# Tuning for {dev_name}",
                f"RDMA_DEV=/sys/class/infiniband/{dev_name}",
                "",
            ]
        )

    lines.extend(
        [
            "# Set recommended sysctl values for RDMA",
            "sysctl -w net.core.rmem_max=16777216 || true",
            "sysctl -w net.core.wmem_max=16777216 || true",
            "sysctl -w net.core.rmem_default=16777216 || true",
            "sysctl -w net.core.wmem_default=16777216 || true",
            "",
            'echo "RDMA tuning complete"',
            "",
        ]
    )

    return "\n".join(lines)


def write_user_stack_configs(
    cfg: Dict,
    snapshot: Dict,
    output_dir: Path,
) -> Dict[str, Path]:
    """Generate and write user stack configuration files.

    Args:
        cfg: lltune configuration dictionary
        snapshot: System snapshot dictionary
        output_dir: Directory to write generated configs

    Returns:
        Dictionary mapping config type to file path
    """
    results = {}
    stack_dir = output_dir / "user_stacks"
    stack_dir.mkdir(parents=True, exist_ok=True)

    # Detect available stacks from snapshot
    user_stack = snapshot.get("user_stack", {})
    onload_version = user_stack.get("onload_version")
    vma_version = user_stack.get("vma_version")

    # Get config settings
    onload_cfg = cfg.get("onload", {}) or {}
    vma_cfg = cfg.get("vma", {}) or {}
    _rdma_cfg = cfg.get("rdma", {}) or {}  # noqa: F841

    # Determine NIC roles
    network_cfg = cfg.get("network", {}) or {}
    interfaces = network_cfg.get("interfaces", []) or []
    trading_ifaces = [
        i.get("name") for i in interfaces if i.get("role") == "trading"
    ]
    primary_iface = trading_ifaces[0] if trading_ifaces else None

    # Generate Onload profile if installed
    if onload_version or onload_cfg.get("generate_profile"):
        tuning = onload_cfg.get("tuning_level", "ultra_low_latency")
        onload_profile = generate_onload_profile(
            nic_role="trading",
            tuning_level=tuning,
            interface=primary_iface,
        )
        onload_path = stack_dir / "lltune-onload.sh"
        onload_path.write_text(onload_profile)
        onload_path.chmod(0o644)
        results["onload_profile"] = onload_path
        logger.info("Generated Onload profile: %s", onload_path)

    # Generate VMA profile if installed
    if vma_version or vma_cfg.get("generate_profile"):
        tuning = vma_cfg.get("tuning_level", "ultra_low_latency")
        vma_profile = generate_vma_profile(
            nic_role="trading",
            tuning_level=tuning,
            interface=primary_iface,
        )
        vma_path = stack_dir / "lltune-vma.sh"
        vma_path.write_text(vma_profile)
        vma_path.chmod(0o644)
        results["vma_profile"] = vma_path
        logger.info("Generated VMA profile: %s", vma_path)

    # Validate and generate RDMA tuning
    rdma_devices = snapshot.get("rdma", {}).get("devices", [])
    nics = snapshot.get("nics", [])

    if rdma_devices:
        # Validate alignment
        issues = validate_rdma_alignment(rdma_devices, nics, cfg)
        for issue in issues:
            if issue.severity == "error":
                logger.error("[%s] %s", issue.component, issue.message)
            elif issue.severity == "warning":
                logger.warning("[%s] %s", issue.component, issue.message)
            else:
                logger.info("[%s] %s", issue.component, issue.message)

        # Generate RDMA tuning script
        rdma_script = generate_rdma_tuning_script(rdma_devices, cfg)
        rdma_path = stack_dir / "lltune-rdma-tune.sh"
        rdma_path.write_text(rdma_script)
        rdma_path.chmod(0o755)
        results["rdma_script"] = rdma_path
        logger.info("Generated RDMA tuning script: %s", rdma_path)

    # Write README
    readme = """\
LLTune User Stack Configurations
================================

This directory contains environment profiles for user-space networking stacks.

Files:
- lltune-onload.sh: Solarflare Onload environment profile
- lltune-vma.sh: Mellanox VMA environment profile
- lltune-rdma-tune.sh: RDMA device tuning script

Usage:

For Onload:
  source /etc/profile.d/lltune-onload.sh
  onload ./your_trading_app

For VMA:
  source /etc/profile.d/lltune-vma.sh
  LD_PRELOAD=libvma.so ./your_trading_app

Installation:
  # Copy profiles to system-wide location
  sudo cp lltune-onload.sh /etc/profile.d/
  sudo cp lltune-vma.sh /etc/profile.d/

Notes:
- These profiles are tuned for ultra-low-latency trading workloads
- Adjust settings based on your specific application requirements
- Monitor CPU usage; spinning modes use 100% CPU on polling threads
- Test thoroughly in a non-production environment first
"""
    readme_path = stack_dir / "README.txt"
    readme_path.write_text(readme)
    results["readme"] = readme_path

    return results
