"""AI and CUDA helpers for training and inference workflows."""

from __future__ import annotations

import math
import os
import platform
import random
import subprocess
from dataclasses import asdict, dataclass, field
from typing import Any, Mapping


@dataclass(frozen=True)
class TrainingPreset:
    """Recommended baseline settings for a training workload."""

    task: str
    accelerator: str
    precision: str
    batch_size: int
    gradient_accumulation_steps: int
    gradient_checkpointing: bool
    compile_model: bool
    dataloader_num_workers: int
    pin_memory: bool
    optimizer: str
    flash_attention: bool

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class TrainingPlan:
    """Actionable training plan built from hardware and task constraints."""

    task: str
    model_scale: str
    accelerator: str
    device: str
    precision: str
    amp_dtype: str
    per_device_batch_size: int
    gradient_accumulation_steps: int
    global_batch_size: int
    gradient_checkpointing: bool
    compile_model: bool
    flash_attention: bool
    dataloader_num_workers: int
    pin_memory: bool
    optimizer: str
    learning_rate: float
    weight_decay: float
    warmup_ratio: float
    max_grad_norm: float
    notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def _load_torch() -> Any | None:
    try:
        import torch  # type: ignore

        return torch
    except Exception:
        return None


def _load_numpy() -> Any | None:
    try:
        import numpy  # type: ignore

        return numpy
    except Exception:
        return None


def _query_nvidia_smi() -> dict[str, Any]:
    try:
        output = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,memory.total,driver_version",
                "--format=csv,noheader,nounits",
            ],
            text=True,
            encoding="utf-8",
            errors="ignore",
        )
    except Exception:
        return {"available": False, "driver_version": None, "gpus": []}

    gpus: list[dict[str, Any]] = []
    driver_version: str | None = None
    for raw_line in output.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        parts = [part.strip() for part in line.split(",")]
        if len(parts) < 3:
            continue
        driver_version = parts[-1]
        memory_mb_str = parts[-2]
        name = ",".join(parts[:-2]).strip()
        try:
            memory_mb = float(memory_mb_str)
        except ValueError:
            memory_mb = 0.0
        gpus.append(
            {
                "name": name,
                "memory_total_gb": round(memory_mb / 1024.0, 2),
                "driver_version": driver_version,
            }
        )

    return {"available": True, "driver_version": driver_version, "gpus": gpus}


def detect_ai_stack() -> dict[str, Any]:
    """Inspect Python/torch/CUDA availability and GPU details."""
    torch = _load_torch()
    smi = _query_nvidia_smi()

    report: dict[str, Any] = {
        "platform": platform.platform(),
        "python_version": platform.python_version(),
        "cpu_logical": os.cpu_count() or 1,
        "torch_installed": bool(torch),
        "torch_version": None,
        "cuda_available": False,
        "cuda_device_count": 0,
        "cuda_devices": [],
        "max_vram_gb": 0.0,
        "cuda_version": None,
        "cudnn_version": None,
        "mps_available": False,
        "nvidia_smi": smi,
    }

    if not torch:
        if smi.get("available") and smi.get("gpus"):
            max_vram = max(float(gpu.get("memory_total_gb", 0.0)) for gpu in smi["gpus"])
            report["cuda_available"] = True
            report["max_vram_gb"] = round(max_vram, 2)
            report["cuda_device_count"] = len(smi["gpus"])
            report["cuda_devices"] = smi["gpus"]
        return report

    report["torch_version"] = getattr(torch, "__version__", None)
    report["cuda_version"] = getattr(getattr(torch, "version", None), "cuda", None)

    cudnn_backend = getattr(getattr(torch, "backends", None), "cudnn", None)
    if cudnn_backend and hasattr(cudnn_backend, "version"):
        try:
            report["cudnn_version"] = cudnn_backend.version()
        except Exception:
            report["cudnn_version"] = None

    mps_backend = getattr(getattr(torch, "backends", None), "mps", None)
    if mps_backend and hasattr(mps_backend, "is_available"):
        try:
            report["mps_available"] = bool(mps_backend.is_available())
        except Exception:
            report["mps_available"] = False

    cuda_devices: list[dict[str, Any]] = []
    cuda_available = False
    try:
        cuda_available = bool(torch.cuda.is_available())
    except Exception:
        cuda_available = False
    report["cuda_available"] = cuda_available

    if cuda_available:
        try:
            device_count = int(torch.cuda.device_count())
        except Exception:
            device_count = 0
        report["cuda_device_count"] = device_count

        for index in range(device_count):
            try:
                props = torch.cuda.get_device_properties(index)
                name = torch.cuda.get_device_name(index)
                capability = f"{props.major}.{props.minor}"
                memory_gb = round(float(props.total_memory) / (1024**3), 2)
            except Exception:
                name = f"cuda:{index}"
                capability = None
                memory_gb = 0.0
            cuda_devices.append(
                {
                    "index": index,
                    "name": name,
                    "capability": capability,
                    "memory_total_gb": memory_gb,
                }
            )
        report["cuda_devices"] = cuda_devices
        if cuda_devices:
            report["max_vram_gb"] = max(
                float(device.get("memory_total_gb", 0.0)) for device in cuda_devices
            )
    elif smi.get("available") and smi.get("gpus"):
        report["cuda_device_count"] = len(smi["gpus"])
        report["cuda_devices"] = smi["gpus"]
        report["max_vram_gb"] = max(float(gpu.get("memory_total_gb", 0.0)) for gpu in smi["gpus"])

    return report


def select_accelerator(hardware: Mapping[str, Any] | None = None) -> str:
    """Choose best available accelerator in this order: CUDA -> MPS -> CPU."""
    details = dict(hardware) if hardware is not None else detect_ai_stack()
    if bool(details.get("cuda_available")):
        return "cuda"
    if bool(details.get("mps_available")):
        return "mps"
    return "cpu"


def set_global_seed(seed: int, deterministic: bool = False) -> dict[str, Any]:
    """Set random seeds for Python, NumPy (if installed), and torch (if installed)."""
    seed_value = int(seed)
    random.seed(seed_value)
    os.environ["PYTHONHASHSEED"] = str(seed_value)

    report: dict[str, Any] = {
        "seed": seed_value,
        "python": True,
        "numpy": False,
        "torch": False,
        "torch_cuda": False,
        "deterministic": bool(deterministic),
    }

    numpy = _load_numpy()
    if numpy is not None:
        numpy.random.seed(seed_value)
        report["numpy"] = True

    torch = _load_torch()
    if torch is None:
        return report

    try:
        torch.manual_seed(seed_value)
        report["torch"] = True
    except Exception:
        report["torch"] = False

    try:
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed_value)
            report["torch_cuda"] = True
    except Exception:
        report["torch_cuda"] = False

    if deterministic:
        try:
            if hasattr(torch, "use_deterministic_algorithms"):
                torch.use_deterministic_algorithms(True)
            cudnn = getattr(getattr(torch, "backends", None), "cudnn", None)
            if cudnn is not None:
                cudnn.deterministic = True
                cudnn.benchmark = False
        except Exception:
            pass
    return report


def tune_torch_runtime(
    *,
    seed: int | None = None,
    deterministic: bool = False,
    benchmark: bool = True,
    allow_tf32: bool = True,
    matmul_precision: str = "high",
    num_threads: int | None = None,
) -> dict[str, Any]:
    """Apply practical PyTorch runtime tweaks for faster training."""
    torch = _load_torch()
    if torch is None:
        raise RuntimeError(
            "PyTorch is not installed. Install torch for your CUDA version, then rerun."
        )

    report: dict[str, Any] = {
        "torch_available": True,
        "seed_applied": False,
        "deterministic": bool(deterministic),
        "benchmark": bool(benchmark),
        "allow_tf32": bool(allow_tf32),
        "matmul_precision": matmul_precision,
        "num_threads": num_threads,
    }

    if seed is not None:
        seed_report = set_global_seed(seed=seed, deterministic=deterministic)
        report["seed_applied"] = True
        report["seed_report"] = seed_report

    cudnn = getattr(getattr(torch, "backends", None), "cudnn", None)
    if cudnn is not None:
        try:
            cudnn.benchmark = bool(benchmark) and not bool(deterministic)
            cudnn.deterministic = bool(deterministic)
            if hasattr(cudnn, "allow_tf32"):
                cudnn.allow_tf32 = bool(allow_tf32)
        except Exception:
            pass

    cuda_backend = getattr(getattr(torch, "backends", None), "cuda", None)
    matmul = getattr(cuda_backend, "matmul", None)
    if matmul is not None and hasattr(matmul, "allow_tf32"):
        try:
            matmul.allow_tf32 = bool(allow_tf32)
        except Exception:
            pass

    if hasattr(torch, "set_float32_matmul_precision"):
        try:
            torch.set_float32_matmul_precision(str(matmul_precision))
        except Exception:
            pass

    if num_threads is not None and hasattr(torch, "set_num_threads"):
        try:
            torch.set_num_threads(max(1, int(num_threads)))
        except Exception:
            pass

    return report


def recommend_training_preset(
    task: str = "general",
    *,
    aggressive: bool = False,
    hardware: Mapping[str, Any] | None = None,
) -> TrainingPreset:
    """Return baseline training settings based on hardware and task type."""
    normalized_task = str(task).strip().lower()
    if normalized_task not in {"general", "cv", "nlp", "llm", "inference"}:
        raise ValueError("task must be one of: general, cv, nlp, llm, inference")

    details = dict(hardware) if hardware is not None else detect_ai_stack()
    accelerator = select_accelerator(details)
    vram_gb = float(details.get("max_vram_gb") or 0.0)
    cpu_logical = int(details.get("cpu_logical") or (os.cpu_count() or 1))

    if accelerator == "cuda":
        precision = "bfloat16" if vram_gb >= 16 else "float16"
    else:
        precision = "float32"

    if normalized_task in {"llm", "nlp"}:
        if vram_gb >= 40:
            batch_size = 8
        elif vram_gb >= 24:
            batch_size = 4
        elif vram_gb >= 12:
            batch_size = 2
        else:
            batch_size = 1
    elif normalized_task == "cv":
        if vram_gb >= 24:
            batch_size = 128
        elif vram_gb >= 12:
            batch_size = 64
        elif vram_gb >= 8:
            batch_size = 32
        else:
            batch_size = 16 if accelerator != "cpu" else 8
    elif normalized_task == "inference":
        batch_size = 1 if aggressive else 8
    else:
        if vram_gb >= 24:
            batch_size = 64
        elif vram_gb >= 12:
            batch_size = 32
        elif vram_gb >= 8:
            batch_size = 16
        else:
            batch_size = 8 if accelerator != "cpu" else 4

    if normalized_task in {"llm", "nlp"}:
        if batch_size <= 1:
            grad_accum = 8 if aggressive else 4
        elif batch_size == 2:
            grad_accum = 2
        else:
            grad_accum = 1
    else:
        grad_accum = 1

    gradient_checkpointing = normalized_task in {"llm", "nlp"} and (
        aggressive or vram_gb < 24
    )
    compile_model = accelerator == "cuda" and normalized_task != "inference"
    dataloader_num_workers = max(2, min(12, cpu_logical // 2))
    pin_memory = accelerator == "cuda"
    optimizer = "adamw_fused" if accelerator == "cuda" else "adamw"
    flash_attention = accelerator == "cuda" and normalized_task in {"llm", "nlp"} and vram_gb >= 12

    return TrainingPreset(
        task=normalized_task,
        accelerator=accelerator,
        precision=precision,
        batch_size=batch_size,
        gradient_accumulation_steps=grad_accum,
        gradient_checkpointing=gradient_checkpointing,
        compile_model=compile_model,
        dataloader_num_workers=dataloader_num_workers,
        pin_memory=pin_memory,
        optimizer=optimizer,
        flash_attention=flash_attention,
    )


def _normalize_model_scale(model_scale: str) -> str:
    normalized = str(model_scale).strip().lower()
    if normalized not in {"tiny", "small", "base", "large", "xl"}:
        raise ValueError("model_scale must be one of: tiny, small, base, large, xl")
    return normalized


def _precision_to_amp_dtype(precision: str) -> str:
    if precision == "bfloat16":
        return "torch.bfloat16"
    if precision == "float16":
        return "torch.float16"
    return "torch.float32"


def build_training_plan(
    task: str = "general",
    *,
    model_scale: str = "base",
    aggressive: bool = False,
    target_global_batch_size: int | None = None,
    hardware: Mapping[str, Any] | None = None,
) -> TrainingPlan:
    """Build an actionable training plan for a given task and hardware profile."""
    details = dict(hardware) if hardware is not None else detect_ai_stack()
    preset = recommend_training_preset(task=task, aggressive=aggressive, hardware=details)
    normalized_scale = _normalize_model_scale(model_scale)

    per_device_batch = max(1, int(preset.batch_size))
    if target_global_batch_size and target_global_batch_size > 0:
        grad_accum = max(1, math.ceil(int(target_global_batch_size) / per_device_batch))
    else:
        grad_accum = max(1, int(preset.gradient_accumulation_steps))
    global_batch = per_device_batch * grad_accum

    task_key = str(task).strip().lower()
    base_lr_by_task = {
        "general": 3e-4,
        "cv": 4e-4,
        "nlp": 2e-4,
        "llm": 2e-4,
        "inference": 1e-5,
    }
    lr_multiplier_by_scale = {
        "tiny": 1.8,
        "small": 1.3,
        "base": 1.0,
        "large": 0.7,
        "xl": 0.5,
    }
    base_lr = base_lr_by_task.get(task_key, 3e-4)
    lr = base_lr * lr_multiplier_by_scale[normalized_scale]
    if aggressive:
        lr *= 0.85

    if task_key in {"llm", "nlp"}:
        weight_decay = 0.1
        warmup_ratio = 0.03
        max_grad_norm = 1.0
    elif task_key == "cv":
        weight_decay = 0.05
        warmup_ratio = 0.05
        max_grad_norm = 1.0
    else:
        weight_decay = 0.01
        warmup_ratio = 0.05
        max_grad_norm = 0.5

    notes: list[str] = []
    if preset.gradient_checkpointing:
        notes.append("Enable gradient checkpointing to reduce peak memory.")
    if preset.flash_attention:
        notes.append("Use Flash Attention kernels if available in your stack.")
    if preset.compile_model:
        notes.append("Try torch.compile() after first stable training run.")
    if target_global_batch_size:
        notes.append("Global batch target enforced via gradient accumulation.")

    return TrainingPlan(
        task=task_key,
        model_scale=normalized_scale,
        accelerator=preset.accelerator,
        device=select_accelerator(details),
        precision=preset.precision,
        amp_dtype=_precision_to_amp_dtype(preset.precision),
        per_device_batch_size=per_device_batch,
        gradient_accumulation_steps=grad_accum,
        global_batch_size=global_batch,
        gradient_checkpointing=preset.gradient_checkpointing,
        compile_model=preset.compile_model,
        flash_attention=preset.flash_attention,
        dataloader_num_workers=preset.dataloader_num_workers,
        pin_memory=preset.pin_memory,
        optimizer=preset.optimizer,
        learning_rate=float(f"{lr:.8g}"),
        weight_decay=weight_decay,
        warmup_ratio=warmup_ratio,
        max_grad_norm=max_grad_norm,
        notes=notes,
    )


def recommend_torch_training_kwargs(plan: TrainingPlan) -> dict[str, Any]:
    """Generate a practical kwargs template for Torch training loops."""
    persistent_workers = bool(plan.dataloader_num_workers > 0 and plan.pin_memory)
    return {
        "device": plan.device,
        "autocast_dtype": plan.amp_dtype,
        "optimizer": {
            "name": plan.optimizer,
            "learning_rate": plan.learning_rate,
            "weight_decay": plan.weight_decay,
        },
        "trainer": {
            "per_device_batch_size": plan.per_device_batch_size,
            "gradient_accumulation_steps": plan.gradient_accumulation_steps,
            "global_batch_size": plan.global_batch_size,
            "gradient_checkpointing": plan.gradient_checkpointing,
            "compile_model": plan.compile_model,
            "max_grad_norm": plan.max_grad_norm,
            "warmup_ratio": plan.warmup_ratio,
        },
        "dataloader": {
            "num_workers": plan.dataloader_num_workers,
            "pin_memory": plan.pin_memory,
            "persistent_workers": persistent_workers,
        },
    }


def apply_cuda_env_defaults(
    *,
    max_split_size_mb: int = 256,
    expandable_segments: bool = True,
    cuda_module_loading: str = "LAZY",
) -> dict[str, str]:
    """Set CUDA-related environment variables commonly used with PyTorch."""
    split_size = max(64, int(max_split_size_mb))
    alloc_value = f"max_split_size_mb:{split_size}"
    if expandable_segments:
        alloc_value += ",expandable_segments:True"

    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = alloc_value
    os.environ["CUDA_MODULE_LOADING"] = str(cuda_module_loading).upper()
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    return {
        "PYTORCH_CUDA_ALLOC_CONF": os.environ["PYTORCH_CUDA_ALLOC_CONF"],
        "CUDA_MODULE_LOADING": os.environ["CUDA_MODULE_LOADING"],
        "TOKENIZERS_PARALLELISM": os.environ["TOKENIZERS_PARALLELISM"],
    }


def optimize_torch_cuda(
    *,
    task: str = "general",
    model_scale: str = "base",
    aggressive: bool = False,
    target_global_batch_size: int | None = None,
    seed: int | None = None,
    deterministic: bool = False,
    benchmark: bool = True,
    allow_tf32: bool = True,
    matmul_precision: str = "high",
    num_threads: int | None = None,
    max_split_size_mb: int = 256,
    expandable_segments: bool = True,
    cuda_module_loading: str = "LAZY",
) -> dict[str, Any]:
    """Apply CUDA env + torch runtime tuning and return a consolidated report."""
    hardware = detect_ai_stack()
    preset = recommend_training_preset(task=task, aggressive=aggressive, hardware=hardware)
    plan = build_training_plan(
        task=task,
        model_scale=model_scale,
        aggressive=aggressive,
        target_global_batch_size=target_global_batch_size,
        hardware=hardware,
    )
    training_kwargs = recommend_torch_training_kwargs(plan)
    env = apply_cuda_env_defaults(
        max_split_size_mb=max_split_size_mb,
        expandable_segments=expandable_segments,
        cuda_module_loading=cuda_module_loading,
    )

    tune_error: str | None = None
    tuned: dict[str, Any] | None = None
    try:
        tuned = tune_torch_runtime(
            seed=seed,
            deterministic=deterministic,
            benchmark=benchmark,
            allow_tf32=allow_tf32,
            matmul_precision=matmul_precision,
            num_threads=num_threads,
        )
    except RuntimeError as error:
        tune_error = str(error)

    return {
        "hardware": hardware,
        "preset": preset.to_dict(),
        "plan": plan.to_dict(),
        "training_kwargs": training_kwargs,
        "env": env,
        "tuned": tuned,
        "tune_error": tune_error,
    }
