import gradio as gr
import outlines
import transformers
import torch

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
pipe = transformers.pipeline("text-generation", "HuggingFaceTB/SmolLM-1.7B-Instruct", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
outlines_tokenizer = outlines.models.TransformerTokenizer(pipe.tokenizer)

### TODO 1: use outliunes with a transformer model made directly
### TODO 2: use a cfg

def string_to_acrostic_grammar(s, dash_initial=True):
    # this will convert a string to a CFG grammar
    chars = filter(str.isalpha, s.upper())
    grammar_rules = [('"- " ' if dash_initial else '') + f'"{char}" /[^-\\r\\n]+/ "\\n"' for char in chars]
    return "?start: " + " ".join(grammar_rules)

def is_this_prompt_a_list(prompt):
    return False
    # ask the model if the prompt is a list, by constraining the generation to yes or no about a question whether the prompt is a list
    question = f'This is a prompt that you have been asked to answer:\n\n```\n{prompt}\n```\n\nIs this prompt asking for a list of items, instead of a story? Begin your answer with "Yes" if asking for a list, otherwise "No", and then give an explanation of why.'
    grammar = '?start: ("Yes" | "No")'
    cfg_logits_processor = outlines.processors.CFGLogitsProcessor(grammar, outlines_tokenizer)
    output = pipe([{"role": "user", "content": question}], logits_processor=transformers.LogitsProcessorList([cfg_logits_processor]), max_new_tokens=10,)
    # output = pipe([{"role": "system", "content": "You are a helpful assistant who answers in one-word answers."}, {"role": "user", "content": question}], max_new_tokens=10,)
    response = output[0]['generated_text'][-1]['content']
    print("is this prompt a list?", response)
    return response == "Yes"


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    acrostic,
    max_tokens,
    temperature,
    top_p,
):
    print({"message": message, "history": history, "system_message": system_message, "acrostic": acrostic, "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p})
    # grammar = "\n".join(['?start: item item item','?item: "- " /[^-\\r\\n]+/ "\\n"'])
    grammar = string_to_acrostic_grammar(acrostic, dash_initial=is_this_prompt_a_list(message))
    two_items_logits_processor = outlines.processors.CFGLogitsProcessor( grammar , outlines_tokenizer )
    output = pipe([{"role": "user", "content": message}], logits_processor=transformers.LogitsProcessorList([two_items_logits_processor]), max_new_tokens=max_tokens,)
    print(output)
    response = output[0]['generated_text'][-1]['content']

    # messages = [{"role": "system", "content": system_message}]

    # for val in history:
    #     if val[0]:
    #         messages.append({"role": "user", "content": val[0]})
    #     if val[1]:
    #         messages.append({"role": "assistant", "content": val[1]})

    # messages.append({"role": "user", "content": message})
    yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Textbox(value="I love you", label="acrostic"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Maximum new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


if __name__ == "__main__":
    demo.launch()