# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 Assem Sabry

"""
GGUF Export Module for ShadeAI.

Provides functionality to convert HuggingFace models (after abliteration)
to GGUF format for use with llama.cpp, Ollama, and other GGUF-compatible
inference engines.

The conversion pipeline:
    [HF Model (abliterated)] -> [save as HF safetensors] -> [convert to GGUF]

This module wraps the official llama.cpp conversion tools and manages
the entire export process automatically.
"""

import os
import sys
import shutil
import subprocess
from pathlib import Path
from typing import Optional

from .utils import print


# Supported GGUF quantization types (ordered by quality, descending)
GGUF_QUANT_TYPES = [
    "f16",       # Full 16-bit float (largest, best quality)
    "q8_0",      # 8-bit quantization (very high quality)
    "q6_k",      # 6-bit k-quant (excellent quality)
    "q5_k_m",    # 5-bit k-quant medium (great quality, recommended)
    "q5_k_s",    # 5-bit k-quant small
    "q4_k_m",    # 4-bit k-quant medium (good balance, most popular)
    "q4_k_s",    # 4-bit k-quant small
    "q3_k_m",    # 3-bit k-quant medium (smallest usable)
    "q3_k_s",    # 3-bit k-quant small
    "q2_k",      # 2-bit k-quant (experimental, significant quality loss)
]

DEFAULT_QUANT_TYPE = "q4_k_m"


def check_llama_cpp() -> Optional[Path]:
    """
    Check if llama.cpp conversion tools are available.
    
    Returns:
        Path to the convert_hf_to_gguf.py script, or None if not found.
    """
    # Check common locations
    search_paths = [
        Path.cwd() / "llama.cpp",
        Path.home() / "llama.cpp",
        Path(os.environ.get("LLAMA_CPP_PATH", "")) if os.environ.get("LLAMA_CPP_PATH") else None,
    ]
    
    for base_path in search_paths:
        if base_path is None:
            continue
        convert_script = base_path / "convert_hf_to_gguf.py"
        if convert_script.exists():
            return convert_script
    
    # Check if llama-cpp-python is installed (has built-in conversion)
    try:
        import llama_cpp  # type: ignore
        return None  # Package exists but we need the script
    except ImportError:
        pass
    
    return None


def clone_llama_cpp(target_dir: Optional[Path] = None) -> Path:
    """
    Clone llama.cpp repository for conversion tools.
    
    Args:
        target_dir: Directory to clone into. Defaults to cwd/llama.cpp.
        
    Returns:
        Path to the cloned llama.cpp directory.
    """
    if target_dir is None:
        target_dir = Path.cwd() / "llama.cpp"
    
    if target_dir.exists():
        print(f"[dim]llama.cpp already exists at {target_dir}[/]")
        return target_dir
    
    print("[bold cyan]Downloading llama.cpp conversion tools...[/]")
    print("[dim]This is a one-time download (~50MB).[/]")
    
    try:
        subprocess.run(
            [
                "git", "clone", "--depth", "1",
                "https://github.com/ggerganov/llama.cpp.git",
                str(target_dir),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        print("[green]✓[/] llama.cpp downloaded successfully.")
        return target_dir
    except FileNotFoundError:
        print("[red]✗ Git is not installed. Please install Git first.[/]")
        print("[dim]Download from: https://git-scm.com/downloads[/]")
        raise RuntimeError("Git is required to download llama.cpp conversion tools.")
    except subprocess.CalledProcessError as e:
        print(f"[red]✗ Failed to clone llama.cpp: {e.stderr}[/]")
        raise RuntimeError("Failed to download llama.cpp.") from e


def install_conversion_deps(llama_cpp_dir: Path) -> None:
    """Install Python dependencies required for GGUF conversion."""
    requirements_file = llama_cpp_dir / "requirements.txt"
    
    if not requirements_file.exists():
        # Try requirements/requirements-convert-hf-to-gguf.txt (newer layout)
        requirements_file = llama_cpp_dir / "requirements" / "requirements-convert-hf-to-gguf.txt"
    
    if requirements_file.exists():
        print("[dim]Installing conversion dependencies...[/]")
        try:
            subprocess.run(
                [sys.executable, "-m", "pip", "install", "-q", "-r", str(requirements_file)],
                check=True,
                capture_output=True,
                text=True,
            )
            print("[green]✓[/] Dependencies installed.")
        except subprocess.CalledProcessError as e:
            print(f"[yellow]Warning: Some dependencies may not have installed: {e.stderr}[/]")
    else:
        # Minimal deps for conversion
        print("[dim]Installing minimal conversion dependencies (gguf, numpy, sentencepiece)...[/]")
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", "gguf", "numpy", "sentencepiece"],
            check=True,
            capture_output=True,
            text=True,
        )


def convert_to_gguf(
    model_path: str | Path,
    output_path: Optional[str | Path] = None,
    quant_type: str = DEFAULT_QUANT_TYPE,
    llama_cpp_dir: Optional[Path] = None,
) -> Path:
    """
    Convert a HuggingFace model to GGUF format.
    
    This is a two-step process:
    1. Convert HF model to GGUF F16 (unquantized GGUF)
    2. Quantize the F16 GGUF to the desired quantization level
    
    Args:
        model_path: Path to the HuggingFace model directory.
        output_path: Path for the output GGUF file. If None, auto-generated.
        quant_type: GGUF quantization type (e.g., "q4_k_m", "q5_k_m", "f16").
        llama_cpp_dir: Path to llama.cpp directory. If None, auto-detected/downloaded.
        
    Returns:
        Path to the output GGUF file.
    """
    model_path = Path(model_path).resolve()
    
    if not model_path.exists():
        raise FileNotFoundError(f"Model directory not found: {model_path}")
    
    quant_type = quant_type.lower()
    if quant_type not in GGUF_QUANT_TYPES:
        raise ValueError(
            f"Unknown quantization type '{quant_type}'. "
            f"Supported types: {', '.join(GGUF_QUANT_TYPES)}"
        )
    
    # Auto-generate output path if not specified
    if output_path is None:
        model_name = model_path.name
        output_path = model_path.parent / f"{model_name}-{quant_type.upper()}.gguf"
    else:
        output_path = Path(output_path)
    
    print()
    print(f"[bold cyan]GGUF Export Pipeline[/]")
    print(f"  Source: [bold]{model_path}[/]")
    print(f"  Output: [bold]{output_path}[/]")
    print(f"  Quantization: [bold]{quant_type.upper()}[/]")
    print()
    
    # Step 1: Find or download llama.cpp
    convert_script = check_llama_cpp()
    if convert_script is None:
        if llama_cpp_dir is None:
            llama_cpp_dir = clone_llama_cpp()
        else:
            llama_cpp_dir = clone_llama_cpp(llama_cpp_dir)
        convert_script = llama_cpp_dir / "convert_hf_to_gguf.py"
        install_conversion_deps(llama_cpp_dir)
    else:
        llama_cpp_dir = convert_script.parent
    
    if not convert_script.exists():
        raise FileNotFoundError(
            f"Conversion script not found at {convert_script}. "
            "Please ensure llama.cpp is properly installed."
        )
    
    # Step 2: Convert HF -> GGUF F16
    f16_output = model_path.parent / f"{model_path.name}-f16.gguf"
    
    print("[bold]Step 1/2:[/] Converting HuggingFace model to GGUF (F16)...")
    try:
        result = subprocess.run(
            [
                sys.executable, str(convert_script),
                str(model_path),
                "--outfile", str(f16_output),
                "--outtype", "f16",
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        print(f"[green]✓[/] F16 GGUF created: {f16_output}")
    except subprocess.CalledProcessError as e:
        print(f"[red]✗ Conversion failed:[/]")
        print(f"[red]{e.stderr}[/]")
        raise RuntimeError("HF to GGUF conversion failed.") from e
    
    # Step 3: Quantize if not F16
    if quant_type == "f16":
        # Just rename the F16 file
        if f16_output != output_path:
            shutil.move(str(f16_output), str(output_path))
        print(f"[green]✓[/] GGUF F16 model saved to: [bold]{output_path}[/]")
        return output_path
    
    print(f"[bold]Step 2/2:[/] Quantizing to {quant_type.upper()}...")
    
    # Find the quantize binary
    quantize_bin = None
    for name in ["llama-quantize", "quantize"]:
        for ext in ["", ".exe"]:
            candidate = llama_cpp_dir / "build" / "bin" / f"{name}{ext}"
            if candidate.exists():
                quantize_bin = candidate
                break
            candidate = llama_cpp_dir / f"{name}{ext}"
            if candidate.exists():
                quantize_bin = candidate
                break
    
    if quantize_bin is None:
        print("[yellow]Warning: llama-quantize binary not found.[/]")
        print("[yellow]The F16 GGUF file has been saved. You can quantize it manually:[/]")
        print(f"  [bold]llama-quantize {f16_output} {output_path} {quant_type.upper()}[/]")
        print()
        print("[dim]To build llama-quantize, run inside the llama.cpp directory:[/]")
        print("  [bold]cmake -B build && cmake --build build --config Release[/]")
        return f16_output
    
    try:
        result = subprocess.run(
            [str(quantize_bin), str(f16_output), str(output_path), quant_type.upper()],
            check=True,
            capture_output=True,
            text=True,
        )
        print(f"[green]✓[/] Quantized GGUF saved to: [bold]{output_path}[/]")
        
        # Clean up the intermediate F16 file
        if f16_output.exists() and f16_output != output_path:
            f16_output.unlink()
            print("[dim]Cleaned up intermediate F16 file.[/]")
        
    except subprocess.CalledProcessError as e:
        print(f"[red]✗ Quantization failed:[/]")
        print(f"[red]{e.stderr}[/]")
        print(f"[yellow]The F16 file is still available at: {f16_output}[/]")
        raise RuntimeError("GGUF quantization failed.") from e
    
    # Print file size comparison
    if output_path.exists():
        size_gb = output_path.stat().st_size / (1024**3)
        print()
        print(f"[bold green]✅ Export complete![/]")
        print(f"  File size: [bold]{size_gb:.2f} GB[/]")
        print()
        print("[bold]Next steps:[/]")
        print(f"  • Run with Ollama:  [bold]ollama create my-model -f Modelfile[/]")
        print(f"  • Run with llama.cpp: [bold]llama-cli -m {output_path}[/]")
    
    return output_path


def register_with_ollama(
    gguf_path: str | Path,
    model_name: Optional[str] = None,
    system_prompt: Optional[str] = None,
) -> bool:
    """
    Register a GGUF model with Ollama for easy local inference.
    
    Args:
        gguf_path: Path to the GGUF model file.
        model_name: Name to register with Ollama. Auto-generated if None.
        system_prompt: Optional system prompt to bake into the Modelfile.
        
    Returns:
        True if registration was successful, False otherwise.
    """
    gguf_path = Path(gguf_path).resolve()
    
    if not gguf_path.exists():
        print(f"[red]✗ GGUF file not found: {gguf_path}[/]")
        return False
    
    if model_name is None:
        model_name = gguf_path.stem.lower().replace(" ", "-") + "-shade"
    
    # Create Modelfile
    modelfile_content = f'FROM "{gguf_path}"\n'
    if system_prompt:
        # Escape quotes in system prompt
        escaped_prompt = system_prompt.replace('"', '\\"')
        modelfile_content += f'\nSYSTEM "{escaped_prompt}"\n'
    
    modelfile_path = gguf_path.parent / "Modelfile"
    with open(modelfile_path, "w", encoding="utf-8") as f:
        f.write(modelfile_content)
    
    print(f"[dim]Registering with Ollama as '{model_name}'...[/]")
    
    try:
        result = subprocess.run(
            ["ollama", "create", model_name, "-f", str(modelfile_path)],
            check=True,
            capture_output=True,
            text=True,
        )
        print(f"[green]✓[/] Model registered with Ollama: [bold]{model_name}[/]")
        print(f"  Run with: [bold]ollama run {model_name}[/]")
        return True
    except FileNotFoundError:
        print("[yellow]Ollama is not installed.[/]")
        print("[dim]Install from: https://ollama.com/download[/]")
        print(f"[dim]Modelfile saved at: {modelfile_path}[/]")
        return False
    except subprocess.CalledProcessError as e:
        print(f"[red]✗ Ollama registration failed: {e.stderr}[/]")
        return False
