import toml
import gradio as gr
from string import Template
from transformers import AutoTokenizer

from vid2persona.gen import tgi_openllm
from vid2persona.gen import local_openllm

tokenizer = None

def _get_system_prompt(
    personality_json_dict: dict,
    prompt_tpl_path: str
) -> str:
    """Assumes a single character is passed."""
    prompt_tpl_path = f"{prompt_tpl_path}/llm.toml"
    system_prompt = Template(toml.load(prompt_tpl_path)['conversation']['system'])

    name = personality_json_dict["name"]
    physcial_description = personality_json_dict["physicalDescription"]
    personality_traits = [str(trait) for trait in personality_json_dict["personalityTraits"]]
    likes = [str(like) for like in personality_json_dict["likes"]]
    dislikes = [str(dislike) for dislike in personality_json_dict["dislikes"]]
    background = [str(info) for info in personality_json_dict["background"]]
    goals = [str(goal) for goal in personality_json_dict["goals"]]
    relationships = [str(relationship) for relationship in personality_json_dict["relationships"]]

    system_prompt = system_prompt.substitute(
        name=name,
        physcial_description=physcial_description,
        personality_traits=', '.join(personality_traits),
        likes=', '.join(likes),
        background=', '.join(background),
        goals=', '.join(goals),
        relationships=', '.join(relationships)
    )

    return system_prompt

async def chat(
    message: str,
    chat_history: list[tuple[str, str]],
    personality_json_dict: dict,
    prompt_tpl_path: str,

    model_id: str,
    max_input_token_length: int,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float,

    hf_token: str,
):
    messages = []
    system_prompt = _get_system_prompt(personality_json_dict, prompt_tpl_path)
    messages.append({"role": "system", "content": system_prompt})
    for user, assistant in chat_history:
        messages.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
    messages.append({"role": "user", "content": message})

    parameters = {
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "repetition_penalty": repetition_penalty
    }

    try:
        for response in local_openllm.send_message(messages, model_id, max_input_token_length, parameters):
            yield response
    except Exception as e:
        gr.Warning(f"{e} ➡️ Switching to TGI remotely hosted model")
    finally:
        async for response in tgi_openllm.send_messages(messages, model_id, hf_token, parameters):
            yield response