import asyncio
import random
from importlib.resources import files
from pathlib import Path
from typing import Union

from litellm import acompletion
from tenacity import retry, wait_exponential, stop_after_attempt

MODEL_RESPONSES_BATCH_SIZE = (
    50  # Size of each batch for model completions. Set this based on API rate limits
)


def generate_random_phrase() -> str:
    """Generate a random phrase combining an adjective and a noun.

    Returns:
        str: A random phrase.
    """

    def get_words_from_file(file_path: Union[str, Path]) -> list[str]:
        with open(file_path, "r") as file:
            return file.read().splitlines()

    adverbs_file_path = files(__package__) / "word-lists/adverbs.txt"
    adverbs = get_words_from_file(adverbs_file_path)
    random_adverb = random.choice(adverbs)
    adjectives_file_path = files(__package__) / "word-lists/adjectives.txt"
    adjectives = get_words_from_file(adjectives_file_path)
    random_adjective = random.choice(adjectives)
    nouns_file_path = files(__package__) / "word-lists/nouns.txt"
    nouns = get_words_from_file(nouns_file_path)
    random_noun = random.choice(nouns)

    return f"{random_adverb}-{random_adjective}-{random_noun}"


async def generate_batch_model_responses(
    model_name: str,
    messages: list[list[dict[str, str]]],
    base_url: str = None,
    api_key: str = None,
) -> list[dict[str, str]]:
    """Generate responses from the specified model for a batch of messages.

    Args:
        model_name (str): The name of the model to generate completions from.
        messages (list[list[dict[str, str]]): A batch of messages to generate completions for.

    Returns:
        list[dict[str, str]]: The generated responses from the model.
    """

    @retry(wait=wait_exponential(multiplier=1, min=2, max=120), stop=stop_after_attempt(5))
    async def generate_model_response(
        model_name: str, messages: list[dict[str, str]]
    ) -> dict[str, str]:
        response = await acompletion(
            model_name, messages, base_url=base_url, api_key=api_key
        )
        return {
            "role": response.choices[0].message.role,
            "content": response.choices[0].message.content,
            "usage": {
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            },
        }

    output = await asyncio.gather(
        *[generate_model_response(model_name, m) for m in messages]
    )
    return output


def generate_model_responses(
    model_name: str,
    messages: list[list[dict[str, str]]],
    system_prompt: str = None,
    base_url: str = None,
    api_key: str = None,
) -> list[dict[str, str]]:
    """
    Generate responses from the specified model for a list of messages. This function batches the messages and
    generates completions for each batch asynchronously.

    Args:
        model_name (str): The name of the model to generate completions from.
        messages (list[list[dict[str, str]]): A list of messages to generate completions for.
        system_prompt (str): The system prompt to use for generating.

    Returns:
        list[dict[str, str]]: The generated responses from the model.
    """
    if system_prompt:
        for message in messages:
            message.insert(0, {"role": "system", "content": system_prompt})
    batches = [
        messages[i : i + MODEL_RESPONSES_BATCH_SIZE]
        for i in range(0, len(messages), MODEL_RESPONSES_BATCH_SIZE)
    ]
    output = []
    for batch in batches:
        batch_output = asyncio.run(
            generate_batch_model_responses(model_name, batch, base_url, api_key)
        )
        output.extend(batch_output)
    return output
