"""Merge LoRA weights and export to GGUF for llama.cpp inference."""

from __future__ import annotations

import argparse
import shutil
from pathlib import Path

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


def export_to_gguf(
    checkpoint_path: Path,
    output_path: Path,
    quantization: str = "Q4_K_M",
) -> Path:
    """Export fine-tuned model to GGUF format.

    Args:
        checkpoint_path: Path to checkpoint directory with LoRA weights
        output_path: Path to save GGUF file
        quantization: Quantization method (Q4_K_M, Q5_K_M, Q8_0)

    Returns:
        Path to exported GGUF file
    """
    print(f"Loading checkpoint from: {checkpoint_path}")

    # Load the base model
    # Note: This assumes the checkpoint was saved with save_pretrained
    # which includes the adapter_config.json

    from unsloth import FastLanguageModel

    # Load model with adapters
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=str(checkpoint_path),
        max_seq_length=2048,
        dtype=None,
        load_in_4bit=False,  # Load full precision for export
    )

    # Merge LoRA weights into base
    print("Merging LoRA weights...")
    model = model.merge_and_unload()

    # Save merged model temporarily
    temp_path = checkpoint_path.parent / "merged"
    temp_path.mkdir(exist_ok=True)

    print(f"Saving merged model to: {temp_path}")
    model.save_pretrained(temp_path)
    tokenizer.save_pretrained(temp_path)

    # Convert to GGUF using llama.cpp
    # Note: This requires llama.cpp's convert script
    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Exporting to GGUF format...")
    print(f"  Quantization: {quantization}")
    print(f"  Output: {output_path}")

    # For now, we'll save in HuggingFace format
    # Full GGUF conversion would require llama.cpp tools
    # which may not be installed in the environment

    # Alternative: Save as merged HF model
    hf_output = output_path.parent / "merged_hf"
    hf_output.mkdir(parents=True, exist_ok=True)

    model.save_pretrained(hf_output)
    tokenizer.save_pretrained(hf_output)

    print(f"\nModel exported to HuggingFace format: {hf_output}")
    print(f"\nTo convert to GGUF, install llama.cpp and run:")
    print(
        f"  python convert_hf_to_gguf.py {hf_output} --outfile {output_path} --outtype {quantization}"
    )

    # Create a marker file
    marker = output_path.parent / "EXPORTED"
    marker.write_text(f"Merged model saved to: {hf_output}\n")

    return hf_output


def merge_only(
    checkpoint_path: Path,
    output_path: Path,
) -> Path:
    """Just merge LoRA weights, save as HF model.

    This is useful if you want to serve via vLLM or HuggingFace directly
    instead of converting to GGUF.
    """
    print(f"Loading checkpoint from: {checkpoint_path}")

    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=str(checkpoint_path),
        max_seq_length=2048,
        dtype=None,
        load_in_4bit=False,
    )

    print("Merging LoRA weights...")
    model = model.merge_and_unload()

    output_path.mkdir(parents=True, exist_ok=True)

    print(f"Saving merged model to: {output_path}")
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)

    print(f"Done! Model saved to: {output_path}")

    return output_path


def main():
    """CLI entry point."""
    parser = argparse.ArgumentParser(description="Export fine-tuned model")
    parser.add_argument(
        "--checkpoint", type=Path, required=True, help="Checkpoint directory"
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("~/.companion/models/exported"),
        help="Output path",
    )
    parser.add_argument("--gguf", action="store_true", help="Export to GGUF format")
    parser.add_argument(
        "--quant", type=str, default="Q4_K_M", help="GGUF quantization type"
    )

    args = parser.parse_args()

    checkpoint = args.checkpoint.expanduser()
    output = args.output.expanduser()

    if args.gguf:
        export_to_gguf(checkpoint, output, args.quant)
    else:
        merge_only(checkpoint, output)


if __name__ == "__main__":
    main()
