"""Preset YAML generator for discovered models.

Generates ready-to-use gpumod preset files from discovered model
metadata and system information.
"""

from __future__ import annotations

import logging
import os
import re
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING

from gpumod.discovery.llamacpp_options import LlamaCppOptions

if TYPE_CHECKING:
    from gpumod.discovery.gguf_metadata import GGUFFile
    from gpumod.discovery.system_info import SystemInfo

logger = logging.getLogger(__name__)

# Template for generated preset YAML
_PRESET_TEMPLATE = """# Generated by gpumod discover on {date}
# Source: {repo_id}
# Quantization: {quant_type}
# Estimated VRAM: {estimated_vram_mb} MB
id: {service_id}
name: {human_name}
driver: llamacpp
port: {port}
vram_mb: {vram_mb}
health_endpoint: /health
startup_timeout: {startup_timeout}
supports_sleep: true
sleep_mode: router
model_id: {model_id}
unit_vars:
  models_dir: $HOME/models
  models_max: 1
  no_models_autoload: true
  default_model: {default_model}
  n_gpu_layers: {n_gpu_layers}
  ctx_size: {ctx_size}
  threads: {threads}
  jinja: true
{extra_vars}"""


@dataclass
class PresetRequest:
    """Request parameters for preset generation.

    Attributes:
        repo_id: HuggingFace repo ID.
        gguf_file: Selected GGUF file metadata.
        system_info: Current system information.
        ctx_size: Desired context window size.
        port: Optional port (auto-assigned if None).
        is_moe: Whether the model is MoE architecture.
    """

    repo_id: str
    gguf_file: GGUFFile
    system_info: SystemInfo
    ctx_size: int = 8192
    port: int | None = None
    is_moe: bool = False


class PresetGenerator:
    """Generates gpumod preset YAML files from model metadata.

    Creates ready-to-use presets that can be added to gpumod with
    'gpumod service add --preset <name>'.

    Example:
        >>> generator = PresetGenerator()
        >>> yaml_str = generator.generate(request)
        >>> print(yaml_str)
    """

    # Characters not allowed in service IDs
    _INVALID_ID_CHARS = re.compile(r"[^a-z0-9-]")

    # Default port range for auto-assignment
    _PORT_RANGE = (7000, 7999)

    # Default startup timeout in seconds
    _DEFAULT_STARTUP_TIMEOUT = 300

    def __init__(
        self,
        *,
        options: LlamaCppOptions | None = None,
    ) -> None:
        """Initialize the generator.

        Args:
            options: LlamaCppOptions instance for config recommendations.
        """
        self._options = options or LlamaCppOptions()

    def generate(self, request: PresetRequest) -> str:
        """Generate preset YAML from request.

        Args:
            request: PresetRequest with model and system info.

        Returns:
            Valid YAML string for a gpumod preset.
        """
        # Derive service ID from model name
        service_id = self._derive_service_id(request.repo_id, request.gguf_file)

        # Extract human-readable name
        human_name = self._derive_human_name(request.repo_id, request.gguf_file)

        # Determine port
        port = request.port or self._find_available_port()

        # Calculate n_gpu_layers based on VRAM budget
        n_gpu_layers = self._calculate_gpu_layers(
            request.gguf_file,
            request.system_info,
            request.is_moe,
        )

        # Get recommended config
        config = self._options.get_recommended_config(
            model_size_b=self._estimate_model_size(request.gguf_file),
            available_vram_mb=request.system_info.gpu_available_mb,
            is_moe=request.is_moe,
            ctx_size=request.ctx_size,
        )

        # Build extra vars for MoE or special cases
        extra_vars = self._build_extra_vars(config, request.is_moe)

        # Determine threads (use half of available CPUs)
        threads = os.cpu_count() or 8
        threads = max(4, threads // 2)

        # Calculate VRAM estimate
        vram_mb = request.gguf_file.estimated_vram_mb
        if n_gpu_layers != -1:
            # Partial offload - reduce VRAM estimate proportionally
            # This is a rough estimate
            vram_mb = int(vram_mb * (n_gpu_layers / 80.0))

        # Get default model name (filename without extension)
        default_model = Path(request.gguf_file.filename).stem

        # Generate YAML
        return _PRESET_TEMPLATE.format(
            date=datetime.now(tz=UTC).strftime("%Y-%m-%d"),
            repo_id=request.repo_id,
            quant_type=request.gguf_file.quant_type or "unknown",
            estimated_vram_mb=request.gguf_file.estimated_vram_mb,
            service_id=service_id,
            human_name=human_name,
            port=port,
            vram_mb=vram_mb,
            startup_timeout=self._DEFAULT_STARTUP_TIMEOUT,
            model_id=request.repo_id,
            default_model=default_model,
            n_gpu_layers=n_gpu_layers,
            ctx_size=request.ctx_size,
            threads=threads,
            extra_vars=extra_vars,
        )

    def _derive_service_id(self, repo_id: str, gguf_file: GGUFFile) -> str:
        """Derive a valid service ID from repo/file info.

        Args:
            repo_id: HuggingFace repo ID.
            gguf_file: GGUF file metadata.

        Returns:
            Valid service ID (lowercase, hyphens only).
        """
        # Start with the repo name part
        name = repo_id.split("/")[-1]

        # Remove common suffixes
        for suffix in ("-GGUF", "-gguf", "_GGUF", "_gguf"):
            name = name.removesuffix(suffix)

        # Convert to lowercase and replace invalid chars
        service_id = name.lower()
        service_id = service_id.replace("_", "-")
        service_id = self._INVALID_ID_CHARS.sub("", service_id)

        # Remove duplicate hyphens
        while "--" in service_id:
            service_id = service_id.replace("--", "-")

        # Trim leading/trailing hyphens
        service_id = service_id.strip("-")

        # Ensure not empty
        if not service_id:
            service_id = "model"

        return service_id

    def _derive_human_name(self, repo_id: str, gguf_file: GGUFFile) -> str:
        """Derive a human-readable name.

        Args:
            repo_id: HuggingFace repo ID.
            gguf_file: GGUF file metadata.

        Returns:
            Human-readable model name.
        """
        name = repo_id.split("/")[-1]

        # Remove GGUF suffix
        for suffix in ("-GGUF", "-gguf", "_GGUF", "_gguf"):
            name = name.removesuffix(suffix)

        # Convert to title case with spaces
        name = name.replace("-", " ").replace("_", " ")

        # Add quantization info
        if gguf_file.quant_type:
            name = f"{name} ({gguf_file.quant_type})"

        return name

    def _find_available_port(self) -> int:
        """Find an available port in the configured range.

        Returns:
            Available port number.
        """
        # For now, just return a default port
        # In a full implementation, would scan existing presets
        return 7070

    def _calculate_gpu_layers(
        self,
        gguf_file: GGUFFile,
        system_info: SystemInfo,
        is_moe: bool,
    ) -> int:
        """Calculate optimal n_gpu_layers for VRAM budget.

        Args:
            gguf_file: GGUF file metadata.
            system_info: Current system information.
            is_moe: Whether model is MoE.

        Returns:
            Number of GPU layers (-1 for all).
        """
        estimated_vram = gguf_file.estimated_vram_mb
        available_vram = system_info.gpu_available_mb

        # Leave 2GB headroom for KV cache and system
        headroom = 2048
        usable_vram = available_vram - headroom

        if estimated_vram <= usable_vram:
            # Full offload fits
            return -1

        # Need partial offload
        # Estimate layers based on VRAM ratio
        # Assume ~80 layers for large models
        estimated_layers = 80
        ratio = usable_vram / estimated_vram

        # Calculate layers, leaving some margin
        n_layers = int(estimated_layers * ratio * 0.85)

        # Ensure at least 1 layer
        return max(1, n_layers)

    def _estimate_model_size(self, gguf_file: GGUFFile) -> float:
        """Estimate model size in billions from file size.

        Args:
            gguf_file: GGUF file metadata.

        Returns:
            Estimated model size in billions of parameters.
        """
        # Very rough estimate: Q4 is ~4 bits per param
        # So file_size_GB * 2 ≈ params in billions
        file_size_gb = gguf_file.size_bytes / (1024**3)
        return file_size_gb * 2

    def _build_extra_vars(
        self,
        config: object,
        is_moe: bool,
    ) -> str:
        """Build extra unit_vars for special cases.

        Args:
            config: RecommendedConfig from options.
            is_moe: Whether model is MoE.

        Returns:
            YAML-formatted extra variables.
        """
        lines: list[str] = []

        if is_moe:
            # Add MoE-specific options
            lines.append("  # MoE model: offload expert computation to CPU")
            lines.append("  n_cpu_moe: 512")

        if hasattr(config, "extra_args"):
            for key, value in config.extra_args.items():
                if key != "n_cpu_moe" or not is_moe:  # Avoid duplicate
                    lines.append(f"  {key}: {value}")

        return "\n".join(lines)


def generate_preset(
    repo_id: str,
    gguf_file: GGUFFile,
    system_info: SystemInfo,
    *,
    ctx_size: int = 8192,
    port: int | None = None,
    is_moe: bool = False,
) -> str:
    """Convenience function to generate a preset.

    Args:
        repo_id: HuggingFace repo ID.
        gguf_file: Selected GGUF file metadata.
        system_info: Current system information.
        ctx_size: Desired context window size.
        port: Optional port (auto-assigned if None).
        is_moe: Whether the model is MoE architecture.

    Returns:
        Valid YAML string for a gpumod preset.
    """
    request = PresetRequest(
        repo_id=repo_id,
        gguf_file=gguf_file,
        system_info=system_info,
        ctx_size=ctx_size,
        port=port,
        is_moe=is_moe,
    )
    generator = PresetGenerator()
    return generator.generate(request)
