import torch
from transformers import AutoTokenizer, LlavaNextProcessor, LlavaNextForConditionalGeneration
from src.logger import get_logger

logger = get_logger(__name__)


def load_model(
    model_name_or_path: str, 
    use_fast: bool = True, 
    torch_dtype=torch.bfloat16
) -> tuple:
    tokenizer, processor = load_tokenizer_and_processor(model_name_or_path, use_fast)
    model = LlavaNextForConditionalGeneration.from_pretrained(
        model_name_or_path, device_map=None, torch_dtype=torch_dtype, attn_implementation="eager"
    )

    return model, tokenizer, processor


def load_tokenizer_and_processor(
    model_name_or_path: str, 
    use_fast: bool = True
) -> tuple:
    tokenizer, processor = None, None
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=use_fast)
        processor = LlavaNextProcessor.from_pretrained(model_name_or_path, use_fast=use_fast)
    except Exception as e:
        logger.warning(
            f"There was an issue loading the tokenizer or processor from {model_name_or_path}. Check if path is correct."
        )

    return tokenizer, processor


def log_trainable_parameters(model):
    def is_lora_param(name: str) -> bool:
        n = name.lower()
        if "lora_" in n or "lokr_" in n or "galore_" in n:
            return True

        return False

    num_total_params = 0
    num_trainable_params = 0
    num_trainable_lora = 0
    num_trainable_base = 0

    trainable_lora_names = []
    trainable_base_names = []

    for name, param in model.named_parameters():
        n = param.numel()
        num_total_params += n
        if param.requires_grad:
            num_trainable_params += n
            if is_lora_param(name):
                num_trainable_lora += n
                trainable_lora_names.append(name)
            else:
                num_trainable_base += n
                trainable_base_names.append(name)

    logger.info(
        (
            f"Trainable params total: {num_trainable_params:,} ({num_trainable_params/num_total_params:.2%}), "
            f"LoRA params: {num_trainable_lora:,} ({num_trainable_lora/num_total_params:.2%}), "
            f"Base params: {num_trainable_base:,} ({num_trainable_base/num_total_params:.2%})"
        )
    )


def get_masking_config(tokenizer) -> tuple:
    """
    Returns a tuple: (assistant_start_ids, response_end_token_id)
    """
    chat_template = tokenizer.chat_template or ""
    
    if "<|start_header_id|>" in chat_template:
        # Llama 3
        start_str = "<|start_header_id|>assistant<|end_header_id|>\n\n"
        start_ids = tokenizer.encode(start_str, add_special_tokens=False)
        
        end_id = tokenizer.convert_tokens_to_ids("<|eot_id|>")
        if end_id is None: 
            end_id = tokenizer.eos_token_id
            
        return start_ids, end_id

    elif "<|im_start|>" in chat_template:
        # ChatML
        start_str = "<|im_start|>assistant\n"
        start_ids = tokenizer.encode(start_str, add_special_tokens=False)
        end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        
        return start_ids, end_id

    elif "[/INST]" in chat_template:
        # Mistral
        start_str = "[/INST]"
        start_ids = tokenizer.encode(start_str, add_special_tokens=False)
        end_id = tokenizer.eos_token_id
        
        return start_ids, end_id

    elif "<start_of_turn>" in chat_template:
        #  Gemma
        start_str = "<start_of_turn>model\n"
        start_ids = tokenizer.encode(start_str, add_special_tokens=False)
        end_id = tokenizer.convert_tokens_to_ids("<end_of_turn>")

        return start_ids, end_id
    else:
        raise ValueError("Unknown chat template format. Please update the function to handle this format.")
