# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import os
import re
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
from peft import PeftModel

def get_device_map():
    num_gpus = torch.cuda.device_count()

    if num_gpus > 1:
        print("More than one GPU found. Setting device_map to use CUDA device 0.")
        return 'cuda:0'
    else:
        return 'auto'

def check_adapter_path(adapters_name):
    """
    Checks if the adapter path is correctly set and not a placeholder.
    Args:
    adapters_name (str): The file path for the adapters.
    Raises:
    ValueError: If the adapters_name contains placeholder characters.
    """
    if '<' in adapters_name or '>' in adapters_name:
        raise ValueError("The adapter path has not been set correctly.")

def load_tokenizer(model_name):
    """
    Loads and returns a tokenizer for the specified model.
    Args:
    model_name (str): The name of the model for which to load the tokenizer.
    Returns:
    AutoTokenizer: The loaded tokenizer with special tokens added and padding side set.
    """
    tok = AutoTokenizer.from_pretrained(model_name, device_map=get_device_map(), trust_remote_code=True)
    tok.add_special_tokens({'pad_token': '[PAD]'})
    tok.padding_side = 'right'  # TRL requires right padding
    return tok

def load_model(model_name, torch_dtype, quant_type):
    """
    Loads and returns a model with the specified quantization configuration.
    If more than one GPU is available, wraps the model with DataParallel.
    Args:
    model_name (str): The name of the model to load.
    torch_dtype (torch.dtype): The data type for model weights (e.g., torch.float16).
    quant_type (str): The quantization type to use.
    Returns:
    AutoModelForCausalLM: The loaded model possibly wrapped with DataParallel.
    """
    try:
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model_name_or_path=model_name,
            trust_remote_code=True,
            device_map=get_device_map(),
            torch_dtype=torch_dtype,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type=quant_type
            ),
        )
      
        return model
    except Exception as e:
        raise RuntimeError(f"Error loading model: {e}")

def resize_embeddings(model, tokenizer):
    """
    Resizes the token embeddings in the model to account for new tokens.
    Args:
    model (AutoModelForCausalLM): The model whose token embeddings will be resized.
    tokenizer (AutoTokenizer): The tokenizer corresponding to the model.
    """
    model.resize_token_embeddings(len(tokenizer))

def load_peft_model(model, adapters_name):
    """
    Loads the PEFT model from the pretrained model and specified adapters.
    Args:
    model (AutoModelForCausalLM): The base model.
    adapters_name (str): Path to the adapters file.
    Returns:
    PeftModel: The PEFT model with the loaded adapters.
    """
    return PeftModel.from_pretrained(model, adapters_name)

def get_device():
    """
    Determines and returns the device to use for computations.
    If CUDA is available, returns a CUDA device, otherwise returns a CPU device.
    Prints the number of GPUs available if CUDA is used.
    Returns:
    torch.device: The device to use.
    """
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Number of GPUs available: {torch.cuda.device_count()}")
    else:
        device = torch.device("cpu")
    return device

def run_prompt(model, tokenizer, device, template):
    """
    Runs an interactive prompt where the user can enter text to get generated responses.
    Continues to prompt the user for input until '#end' is entered.
    Args:
    model (AutoModelForCausalLM): The model to use for text generation.
    tokenizer (AutoTokenizer): The tokenizer to use for encoding the input text.
    device (torch.device): The device on which to perform the computation.
    template (str): The template string to format the input text.
    """
    while True:
        new_input = input("Enter your text (type #end to stop): ")
        if new_input == "#end":
            break

        try:
            _ = generate_text(model, tokenizer, device, new_input, template)
        except Exception as e:
            print(f"An error occurred during text generation: {e}")
            
def generate_text(model, tokenizer, device, input_text, template):
    """
    Generates and returns text using the provided model and tokenizer for the input text.
    Args:
    model (AutoModelForCausalLM): The model to use for text generation.
    tokenizer (AutoTokenizer): The tokenizer to use for encoding the input text.
    device (torch.device): The device on which to perform the computation.
    input_text (str): The input text to generate responses for.
    template (str): The template string to format the input text.
    Returns:
    torch.Tensor: The generated text tensor.
    """
    inputs = tokenizer(template.format(input_text), return_tensors="pt")
    inputs = inputs.to(device)  # Move input tensors to the device
    streamer = TextStreamer(tokenizer)
    return model.generate(**inputs, streamer=streamer,
                          max_new_tokens=1024,
                          pad_token_id=tokenizer.pad_token_id,
                          eos_token_id=tokenizer.eos_token_id)

def get_last_folder_alphabetically(directory_path):
    """
    Finds the last folder alphabetically in a specified directory.

    Args:
        directory_path (str): The path to the directory.

    Returns:
        str: The path to the last folder found alphabetically.
        If the directory does not exist or contains no folders, a descriptive string is returned.
    """
    if not os.path.exists(directory_path):
        return "Directory does not exist."

    all_files_and_folders = os.listdir(directory_path)
    only_folders = [f for f in all_files_and_folders if os.path.isdir(os.path.join(directory_path, f))]
    if not only_folders:
        return "No folders found in the directory."

    only_folders.sort(key=natural_sort_key)
    last_folder = only_folders[-1]
    return os.path.join(directory_path, last_folder)

def natural_sort_key(s):
    """
    Generates a key for sorting strings that contain numbers where the numbers should be sorted numerically,
    and the rest alphabetically.

    Args:
        s (str): The string to be sorted.

    Returns:
        list: A list of strings and integers derived from the input string.
    """
    return [int(text) if text.isdigit() else text.lower() for text in re.split('([0-9]+)', s)]