import os
import torch
import logging
import gc
import json
from pathlib import Path
from unsloth import FastVisionModel
from safetensors.torch import save_file

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

INPUT_MODEL = "punjabi_gemma/ankahi"
OUTPUT_DIR = "artifacts/deploy/ankahi-gemma4-e4b-int8"

def main():
    log.info(f"Loading merged BF16 model from {INPUT_MODEL} for quantization...")
    
    # Reload with INT8
    model_int8, processor = FastVisionModel.from_pretrained(
        INPUT_MODEL,
        load_in_4bit=False,  # we want 8bit
        load_in_8bit=True,   # unsloth handles this via bnb
        device_map="auto",
        dtype=torch.bfloat16,
    )

    log.info(f"Saving INT8 model to {OUTPUT_DIR}...")
    Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
    
    log.info("  Performing manual Safetensors save to avoid serialization issues...")
    
    # 1. Save weights manually using safetensors
    state_dict = model_int8.state_dict()
    
    # Handle shared tensors: model.language_model.embed_tokens.weight and lm_head.weight
    shared_keys = [
        ("model.language_model.embed_tokens.weight", "lm_head.weight"),
        ("model.embed_tokens.weight", "lm_head.weight")
    ]
    
    for k1, k2 in shared_keys:
        if k1 in state_dict and k2 in state_dict:
            if state_dict[k1].data_ptr() == state_dict[k2].data_ptr():
                log.info(f"  Detected shared weights ({k1} and {k2}). Removing {k2} for safetensors save.")
                del state_dict[k2]
                break

    save_file(state_dict, os.path.join(OUTPUT_DIR, "model.safetensors"))
    
    # 2. Manually fix config.json by copying from INPUT_MODEL and adding quantization info
    with open(os.path.join(INPUT_MODEL, "config.json"), "r") as f:
        clean_config = json.load(f)
    
    # Scrub non-serializable or problematic fields if any
    # (Just in case, though the file on disk should be clean)
    
    clean_config["quantization_config"] = {
        "bits": 8,
        "quant_method": "bitsandbytes",
        "load_in_8bit": True
    }
    clean_config["_name_or_path"] = "ankahi-gemma4-e4b-int8"
    
    with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
        json.dump(clean_config, f, indent=2)
    
    # 3. Save processor
    processor.save_pretrained(OUTPUT_DIR)
    
    # 4. Copy other necessary files
    for filename in ["tokenizer.json", "tokenizer_config.json", "generation_config.json", "chat_template.jinja", "processor_config.json"]:
        src = os.path.join(INPUT_MODEL, filename)
        if os.path.exists(src):
            import shutil
            shutil.copy(src, os.path.join(OUTPUT_DIR, filename))

    log.info(f"Manual INT8 save complete in {OUTPUT_DIR}")

if __name__ == "__main__":
    main()