"""VRAM-based optimization rules (P0-P1)."""

import math
import re

from bilgeai.analyzer.models import DatasetProfile
from bilgeai.core.constants import VRAM_LOW, VRAM_MEDIUM, VRAM_VERY_LOW
from bilgeai.optimizer.models import TrainingConfig
from bilgeai.optimizer.rules.base import Rule
from bilgeai.profiler.models import HardwareProfile


def _guess_model_size_b(model_name: str) -> float:
    """Guess model size in billions from model name."""
    name_lower = model_name.lower()
    match = re.search(r"(\d+\.?\d*)b", name_lower)
    if match:
        return float(match.group(1))
    for size_str in ["70", "34", "13", "8", "7", "3", "1.5", "0.5"]:
        if size_str in name_lower:
            return float(size_str)
    return 7.0


class VRAMQuantizationRule(Rule):
    """Decide quantization based on per-GPU VRAM (bottleneck in data parallel)."""

    name = "vram_quantization"
    priority = 0

    def apply(
        self,
        config: TrainingConfig,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
    ) -> TrainingConfig:
        if not hardware.gpus:
            config.quantization = "4bit"
            config.add_explanation(
                decision='quantization = "4bit"',
                reason="No GPU detected. 4-bit quantization for CPU/minimal training.",
                confidence=0.85,
                alternatives=["Use a machine with GPU for faster training"],
            )
            return config

        # In data parallel, each GPU must fit the full model — use min VRAM
        vram = hardware.min_vram_gb
        gpu_note = ""
        if hardware.num_gpus > 1:
            gpu_note = f" (smallest of {hardware.num_gpus} GPUs)"

        if vram < VRAM_VERY_LOW:
            config.quantization = "4bit"
            config.use_gradient_checkpointing = True
            config.add_explanation(
                decision='quantization = "4bit"',
                reason=f"GPU VRAM ({vram:.0f} GB{gpu_note}) is very low. 4-bit required.",
                confidence=0.95,
                alternatives=['quantization = "8bit" (may OOM)'],
            )
        elif vram < VRAM_LOW:
            config.quantization = "4bit"
            config.add_explanation(
                decision='quantization = "4bit"',
                reason=f"GPU VRAM ({vram:.0f} GB{gpu_note}) < {VRAM_LOW} GB. 4-bit recommended.",
                confidence=0.90,
                alternatives=['quantization = "8bit" (feasible with small batch)'],
            )
        elif vram < VRAM_MEDIUM:
            config.quantization = "4bit"
            config.add_explanation(
                decision='quantization = "4bit"',
                reason=f"GPU VRAM ({vram:.0f} GB{gpu_note}). 4-bit offers best efficiency.",
                confidence=0.85,
                alternatives=['quantization = "8bit"', 'quantization = "none" (risky)'],
            )
        else:
            config.quantization = "none"
            config.add_explanation(
                decision='quantization = "none"',
                reason=(
                    f"GPU VRAM ({vram:.0f} GB{gpu_note}) >= {VRAM_MEDIUM} GB. "
                    "Full precision feasible."
                ),
                confidence=0.80,
                alternatives=['quantization = "4bit" (saves memory)', 'quantization = "8bit"'],
            )

        return config


class VRAMBatchSizeRule(Rule):
    """Decide per-GPU batch size with conservative VRAM safety margins."""

    name = "vram_batch_size"
    priority = 1

    def apply(
        self,
        config: TrainingConfig,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
    ) -> TrainingConfig:
        if not hardware.gpus:
            config.batch_size = 1
            config.gradient_accumulation_steps = 32
            config.add_explanation(
                decision="batch_size = 1, gradient_accumulation_steps = 32",
                reason="No GPU. Minimal batch with gradient accumulation.",
                confidence=0.90,
            )
            return config

        vram = hardware.min_vram_gb
        num_gpus = hardware.num_gpus

        quant_factor = {"4bit": 0.25, "8bit": 0.5, "none": 1.0}.get(
            config.quantization or "4bit", 0.25
        )

        # Model weight memory (scales with actual model size)
        model_b = _guess_model_size_b(config.model_name)
        base_model_mem = model_b * 2.0 * quant_factor

        # LoRA adapter + optimizer states
        lora_overhead = 1.5 if config.lora_enabled else 0.0

        # Safety headroom: CUDA kernels, fragmentation, etc.
        headroom = 2.0
        available = vram - base_model_mem - lora_overhead - headroom

        # Per-sample activation memory scales with seq_length and model size
        seq_factor = config.max_seq_length / 512.0
        per_sample = max(0.3, 0.8 * (model_b / 7.0) * seq_factor * quant_factor)

        if available <= 0:
            config.batch_size = 1
        else:
            config.batch_size = max(1, min(16, math.floor(available / per_sample)))

        # Hard cap for small VRAM GPUs
        if vram < 10:
            config.batch_size = min(config.batch_size, 4)
        elif vram < 16:
            config.batch_size = min(config.batch_size, 8)

        # Effective batch size target: 32 (stable for most SFT)
        effective_target = 32
        effective_per_step = config.batch_size * num_gpus
        if effective_per_step < effective_target:
            config.gradient_accumulation_steps = max(
                1, min(64, effective_target // effective_per_step)
            )

        effective = config.batch_size * num_gpus * config.gradient_accumulation_steps

        gpu_info = f"VRAM {vram:.0f} GB"
        if num_gpus > 1:
            gpu_info += f" x {num_gpus} GPUs"

        config.add_explanation(
            decision=(
                f"batch_size = {config.batch_size}, "
                f"gradient_accumulation = {config.gradient_accumulation_steps}"
            ),
            reason=(
                f"{gpu_info}. Model ~{base_model_mem:.1f} GB ({config.quantization}), "
                f"LoRA ~{lora_overhead:.1f} GB, headroom {headroom:.0f} GB. "
                f"Available ~{max(0, available):.1f} GB. "
                f"Effective batch = {effective}."
            ),
            confidence=0.85,
            alternatives=[
                f"batch_size = {max(1, config.batch_size // 2)} (safer, less OOM risk)",
            ],
        )

        return config


class DistributedStrategyRule(Rule):
    """Decide distributed training strategy based on GPU count and VRAM."""

    name = "distributed_strategy"
    priority = 1

    def apply(
        self,
        config: TrainingConfig,
        hardware: HardwareProfile,
        dataset: DatasetProfile,
    ) -> TrainingConfig:
        num_gpus = hardware.num_gpus
        config.num_gpus = max(1, num_gpus)
        config.num_nodes = 1  # Local detection is single-node

        if num_gpus <= 1:
            config.distributed_strategy = "none"
            return config

        # Multi-GPU: choose strategy
        vram = hardware.min_vram_gb

        if vram >= VRAM_MEDIUM and config.quantization == "none":
            # Large GPUs, full precision → FSDP shards model across GPUs
            config.distributed_strategy = "fsdp"
            config.add_explanation(
                decision=f"distributed_strategy = fsdp ({num_gpus} GPUs)",
                reason=(
                    f"{num_gpus} GPUs with {vram:.0f} GB each. "
                    "FSDP shards model + optimizer across GPUs for memory efficiency."
                ),
                confidence=0.85,
                alternatives=["ddp (simpler, but each GPU holds full model)"],
            )
        else:
            # LoRA + quantization → DDP is simpler and works well
            config.distributed_strategy = "ddp"
            config.add_explanation(
                decision=f"distributed_strategy = ddp ({num_gpus} GPUs)",
                reason=(
                    f"{num_gpus} GPUs detected. DDP replicates model on each GPU. "
                    "With LoRA/quantization, per-GPU memory is already small."
                ),
                confidence=0.90,
                alternatives=[
                    "fsdp (more complex, useful for full fine-tuning)",
                    "deepspeed (ZeRO stages, advanced memory optimization)",
                ],
            )

        return config
