import os
from global_var import get_model, get_tokenizer

def save_model_to_dir(save_model_name, models_dir, model, tokenizer, save_method="default"):
    """
    保存模型到指定目录
    :param save_model_name: 要保存的模型名称
    :param models_dir: 模型保存的基础目录
    :param model: 要保存的模型
    :param tokenizer: 要保存的tokenizer
    :param save_method: 保存模式选项
        - "default": 默认保存方式
        - "merged_16bit": 合并为16位
        - "merged_4bit": 合并为4位
        - "lora": 仅LoRA适配器
        - "gguf": 保存为GGUF格式
        - "gguf_q4_k_m": 保存为q4_k_m GGUF格式
        - "gguf_f16": 保存为16位GGUF格式
    :return: 保存结果消息或错误信息
    """
    try:
        if model is None:
            return "没有加载的模型可保存"
            
        save_path = os.path.join(models_dir, save_model_name)
        os.makedirs(save_path, exist_ok=True)
        
        if save_method == "default":
            model.save_pretrained(save_path)
            tokenizer.save_pretrained(save_path)
        elif save_method == "merged_16bit":
            model.save_pretrained_merged(save_path, tokenizer, save_method="merged_16bit")
        elif save_method == "merged_4bit":
            model.save_pretrained_merged(save_path, tokenizer, save_method="merged_4bit_forced")
        elif save_method == "lora":
            model.save_pretrained_merged(save_path, tokenizer, save_method="lora")
        elif save_method == "gguf":
            model.save_pretrained_gguf(save_path, tokenizer)
        elif save_method == "gguf_q4_k_m":
            model.save_pretrained_gguf(save_path, tokenizer, quantization_method="q4_k_m")
        elif save_method == "gguf_f16":
            model.save_pretrained_gguf(save_path, tokenizer, quantization_method="f16")
        else:
            return f"不支持的保存模式: {save_method}"
            
        return f"模型已保存到 {save_path} (模式: {save_method})"
    except Exception as e:
        return f"保存模型时出错: {str(e)}"