import argparse
import os
from transformers.trainer_utils import get_last_checkpoint
from peft import PeftConfig, PeftModel
from src.logger import get_logger
from src.modeling.llava.utils import load_model


logger = get_logger(__name__)


def _is_peft_checkpoint(path: str) -> bool:
    try:
        PeftConfig.from_pretrained(path)
        return True
    except Exception:
        return False


def merge_and_save(base_model_name_or_path: str, trainer_output_dir: str, final_model_name_or_path: str):
    last_checkpoint = get_last_checkpoint(trainer_output_dir)
    if last_checkpoint is None:
        raise ValueError(f"No checkpoint found in {trainer_output_dir}")

    logger.info(f"Using checkpoint: {last_checkpoint}")

    os.makedirs(final_model_name_or_path, exist_ok=True)

    if _is_peft_checkpoint(last_checkpoint):
        logger.info("Detected PEFT checkpoint.")
        logger.info(f"Loading base model: {base_model_name_or_path}")
        model, tokenizer, processor = load_model(base_model_name_or_path)
        model.to("cpu").eval()
        
        logger.info("Attaching LoRA adapter and merging into base (merge_and_unload).")
        model = PeftModel.from_pretrained(model, last_checkpoint, is_trainable=False)
        model = model.merge_and_unload()
    else:
        logger.info("Assuming full model checkpoint.")
        model, tokenizer, processor = load_model(last_checkpoint)
        model.to("cpu").eval()

        if not tokenizer or not processor:
            logger.warning("Tokenizer or processor could not be loaded from the checkpoint. Attempting to load from base model.")
            _, tokenizer, processor = load_model(base_model_name_or_path)

    logger.info(f"Saving merged model to: {final_model_name_or_path}")
    model.save_pretrained(final_model_name_or_path, safe_serialization=True)
    tokenizer.save_pretrained(final_model_name_or_path)
    processor.save_pretrained(final_model_name_or_path)

    logger.info("Done.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Merge LoRA adapter into base LLaVA model (adapter-free export).")
    parser.add_argument("--base_model_name_or_path", required=True)
    parser.add_argument("--trainer_output_dir", required=True)
    parser.add_argument("--final_model_name_or_path", required=True)
    args = parser.parse_args()

    if not os.path.isdir(args.trainer_output_dir):
        raise FileNotFoundError(f"Run directory not found: {args.trainer_output_dir}")
    
    merge_and_save(args.base_model_name_or_path, args.trainer_output_dir, args.final_model_name_or_path)
