"""
Merge the trained LoRA adapter into the Phi-3 base model and save a standalone
fp16 model. This merged model is the input to GGUF conversion.

Run on your laptop (the base model is already cached there):

    python scripts/merge_adapter.py

Env overrides (optional):
    BASE_MODEL_ID   default: microsoft/Phi-3-mini-4k-instruct
    ADAPTER_PATH    default: models/phi3-text-to-sql-adapter
    MERGED_OUT      default: models/phi3-text-to-sql-merged

Notes:
- Merging is weight arithmetic (W' = W + (alpha/r) * B@A); it is mathematically
  equivalent to applying the adapter. Under greedy decoding the SQL output is the
  same as the current fine-tuned model.
- fp16 keeps the merged checkpoint ~7.6 GB. GGUF conversion + Q4_K_M quantization
  shrinks it to ~2.3 GB.
"""
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL_ID = os.environ.get("BASE_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
ADAPTER_PATH = os.environ.get("ADAPTER_PATH", "models/phi3-text-to-sql-adapter")
MERGED_OUT = os.environ.get("MERGED_OUT", "models/phi3-text-to-sql-merged")


def main():
    print(f"Base model : {BASE_MODEL_ID}")
    print(f"Adapter    : {ADAPTER_PATH}")
    print(f"Output     : {MERGED_OUT}")

    print("\nLoading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=False)

    print("Loading base model in fp16 (CPU)...")
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.float16,
        device_map=None,
        low_cpu_mem_usage=True,
        trust_remote_code=False,
    )

    print("Attaching adapter...")
    model = PeftModel.from_pretrained(model, ADAPTER_PATH)

    print("Merging adapter into base weights (merge_and_unload)...")
    model = model.merge_and_unload()

    os.makedirs(MERGED_OUT, exist_ok=True)
    print(f"Saving merged model to {MERGED_OUT} ...")
    model.save_pretrained(MERGED_OUT, safe_serialization=True)
    tokenizer.save_pretrained(MERGED_OUT)

    # Phi-3's GGUF converter needs the original SentencePiece tokenizer.model,
    # which save_pretrained() does not emit for the fast tokenizer. Copy it over.
    tok_model_dst = os.path.join(MERGED_OUT, "tokenizer.model")
    if not os.path.exists(tok_model_dst):
        try:
            import shutil
            from huggingface_hub import hf_hub_download
            src = hf_hub_download(BASE_MODEL_ID, "tokenizer.model")
            shutil.copy(src, tok_model_dst)
            print("Copied tokenizer.model (required by the GGUF converter).")
        except Exception as e:
            print(f"WARNING: could not fetch tokenizer.model ({e}). "
                  "Copy it manually from the base model before converting.")

    print("\nDone. Next: convert to GGUF (see scripts/CONVERT_GGUF.md).")


if __name__ == "__main__":
    main()