import logging
from huggingface_hub import HfApi, create_repo, upload_folder
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
from transformers.models.llama.modeling_llama import LlamaConfig, LlamaForCausalLM
from torch import nn

# Importar BitLinear desde el módulo de cuantización existente
from src.training.utils.quantization import BitLinear
from src.training.utils.linear_to_bitlinear import replace_linears_in_hf

# Cuantización y conversión a BitNet

class BitNetConfig(LlamaConfig):
    model_type = "bitnet"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

class BitNetModel(PreTrainedModel):
    config_class = BitNetConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaForCausalLM(config)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

# Registrar BitNet como un nuevo tipo de modelo en Hugging Face Transformers
AutoConfig.register("bitnet", BitNetConfig)
AutoModelForCausalLM.register(BitNetConfig, BitNetModel)

def save_convert_and_push_model(trainer, output_path, huggingface_id, new_model_name, hf_token, convert_to_bitnet=False, do_push=True):
    """
    Convierte el modelo a BitNet si se solicita, guarda el modelo y opcionalmente lo sube a Hugging Face Hub.

    Args:
        trainer: Trainer instance con el modelo ya entrenado.
        output_path: Ruta para guardar el modelo localmente.
        huggingface_id: ID del usuario u organización en Hugging Face Hub.
        new_model_name: Nombre del nuevo modelo que se subirá al Hub.
        hf_token: Token de acceso para Hugging Face Hub.
        convert_to_bitnet (bool): Flag para decidir si se convierte a BitNet.
        do_push (bool): Flag para decidir si se hace push al Hugging Face Hub.

    Returns:
        None
    """
    logger = logging.getLogger(__name__)

    # Convertir a BitNet si se especifica
    if convert_to_bitnet:
        logger.info("🔄 Converting the model to BitNet architecture.")
        replace_linears_in_hf(trainer.model)
        output_dir = f"{output_path}/bitnet_model"
    else:
        output_dir = f"{output_path}/final_model"

    # Guardar el modelo localmente
    logger.info(f"💾 Saving the model locally to {output_dir}.")
    trainer.model.save_pretrained(output_dir)
    trainer.tokenizer.save_pretrained(output_dir)

    # Subir el modelo al Hugging Face Hub si se solicita
    if do_push:
        logger.info("☁️ Uploading the model and tokenizer to Hugging Face Hub.")
        api = HfApi()
        create_repo(
            repo_id=f"{huggingface_id}/{new_model_name}",
            repo_type="model",
            exist_ok=True,
            token=hf_token,
        )
        upload_folder(
            folder_path=output_dir,
            repo_type="model",
            repo_id=f"{huggingface_id}/{new_model_name}",
            token=hf_token,
        )