import gradio as gr
from transformers import AutoTokenizer
import json
import os
from huggingface_hub import login

# Fetch HF Token
HUGGINGFACEHUB_API_TOKEN = os.environ.get("HF_TOKEN", "")

default_model = "meta-llama/Meta-Llama-3-8B-Instruct"

demo_conversation = """[
    {"role": "system", "content": "You are a helpful chatbot."},
    {"role": "user", "content": "Hi there!"},
    {"role": "assistant", "content": "Hello, human!"},
    {"role": "user", "content": "Can I ask a question?"}
]"""

description_text = """# Chat Template Viewer
### This space is a helper to learn more about [Chat Templates](https://huggingface.co/docs/transformers/main/en/chat_templating).
"""

default_tools = [{"type": "function", "function": {"name": "get_current_weather",
                "description": "Get the current weather",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"},
                        "format": {"type": "string", "enum": ["celsius", "fahrenheit"],
                                   "description": "The temperature unit to use. Infer this from the user's location."}
                    },
                    "required": ["location", "format"]
                }
            }}]

def get_template_names(model_name):
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if isinstance(tokenizer.chat_template, dict):
            return list(tokenizer.chat_template.keys())
        else:
            return []
    except Exception as e:
        return ["Default"]

def update_template_dropdown(model_name):
    template_names = get_template_names(model_name)
    if template_names:
        return gr.Dropdown.update(choices=template_names, value=template_names[0])
    return gr.Dropdown.update(choices=[], value=None)

def apply_chat_template(model_name, test_conversation, add_generation_prompt, cleanup_whitespace, template_name, hf_token, kwargs):
    try:
        if hf_token:
            login(token=hf_token)  # Ensure login is successful

        tokenizer = AutoTokenizer.from_pretrained(model_name)
    except Exception as e:
        return f"Error: Could not load model {model_name} or invalid HF token. {str(e)}"

    try:
        conversation = json.loads(test_conversation)
        
        if template_name and tokenizer.chat_template:
            template = tokenizer.chat_template.get(template_name, None)
        else:
            template = None

        formatted = tokenizer.apply_chat_template(
            conversation,
            chat_template=template,
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
            tools=default_tools
        )
        return formatted
    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown(description_text)

    model_name_input = gr.Textbox(label="Model Name", placeholder="Enter model name", value=default_model)
    template_dropdown = gr.Dropdown(label="Template Name", choices=[], interactive=True)
    conversation_input = gr.TextArea(value=demo_conversation, lines=6, label="Conversation")
    add_generation_prompt_checkbox = gr.Checkbox(value=False, label="Add generation prompt")
    cleanup_whitespace_checkbox = gr.Checkbox(value=True, label="Cleanup template whitespace")
    hf_token_input = gr.Textbox(label="Hugging Face Token (optional)", placeholder="Enter your HF token", type="password")
    kwargs_input = gr.JSON(label="Additional kwargs", value=default_tools, visible=False)
    output = gr.TextArea(label="Formatted conversation", interactive=False)

    update_button = gr.Button("Update Template List")
    format_button = gr.Button("Format Conversation")

    update_button.click(fn=update_template_dropdown, inputs=model_name_input, outputs=template_dropdown)

    format_button.click(
        fn=apply_chat_template,
        inputs=[
            model_name_input,
            conversation_input,
            add_generation_prompt_checkbox,
            cleanup_whitespace_checkbox,
            template_dropdown,
            hf_token_input,
            kwargs_input
        ],
        outputs=output
    )

demo.launch()