import gradio as gr
import torch
import os

from model import get_input_token_length, run

DEFAULT_SYSTEM_PROMPT = """\
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
"""
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = 4000


LICENSE = """
<p/>

---
As a derivate work of [Llama-2-13b-chat](https://huggingface.co/meta-llama/Llama-2-13b-chat) by Meta,
this demo is governed by the original [license](https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat/blob/main/LICENSE.txt) and [acceptable use policy](https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat/blob/main/USE_POLICY.md).
"""

is_spaces = True if "SPACE_ID" in os.environ else False
if is_spaces :
    is_shared_ui = True if "gradio-discord-bots/llama-2-13b-chat-transformers" in os.environ['SPACE_ID'] else False
else:
    is_shared_ui = False
is_gpu_associated = torch.cuda.is_available()


def generate(
    message: str,
    history_with_input: list[tuple[str, str]],
    system_prompt=DEFAULT_SYSTEM_PROMPT,
    max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
    temperature=1.0,
    top_p=0.95,
    top_k=50,
) -> tuple[str, list[tuple[str, str]]]:
    if is_shared_ui:
        raise ValueError("Cannot use demo running in shared_ui. Must duplicate your own space.")
    if max_new_tokens > MAX_MAX_NEW_TOKENS:
        raise ValueError
    
    history = history_with_input[:-1]
    input_token_length = get_input_token_length(message, history, system_prompt)
    if input_token_length > MAX_INPUT_TOKEN_LENGTH:
        response = f'The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Please create a new thread.'
    else:
        response = run(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k)
    return response, history + [(message, response)]


with gr.Blocks() as demo:

    gr.Markdown(
        """
    # Llama-2-13b-chat-hf Discord Bot Powered by Gradio and Hugging Face Transformers
    
    ### First install the `gradio_client`
        
    ```bash
    pip install gradio_client
    ```
    
    ### Then deploy to discord in one line! ⚡️
    
    ```python
    secrets = {"HUGGING_FACE_HUB_TOKEN": "<your-key-here>",}
    client = grc.Client.duplicate("gradio-discord-bots/llama-2-13b-chat-transformers", secrets=secrets, hardware="a10g-small")
    client.deploy_discord(api_names=["chat"])
    ```
    """
    )

    gr.Markdown(LICENSE)
    with gr.Row(visible=False):
        state = gr.State([])
        msg = gr.Textbox()
        output = gr.Textbox()
        btn = gr.Button()
        btn.click(generate, [msg, state], [output, state], api_name="chat")

demo.queue(max_size=20).launch()