import gradio as gr
from huggingface_hub import InferenceClient

# Define available models and their Hugging Face IDs
available_models = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    # Add more models here as needed
}

def respond(
    message: str,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    model_name: str,
):
    """
    Generates a response from the AI model based on the user's message and chat history.
    """
    client = InferenceClient(model=available_models[model_name])

    messages = [{"role": "system", "content": system_message}]
    for user_input, assistant_response in history:
        messages.append({"role": "user", "content": user_input})
        messages.append({"role": "assistant", "content": assistant_response})
    messages.append({"role": "user", "content": message})

    streamed_response = ""
    try:
        for response in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            chunk = response.choices[0].delta.get("content", "")
            streamed_response += chunk
            yield streamed_response
    except Exception as e:
        yield f"**Error:** {str(e)}"

# Latest updates
latest_updates = """
**Chatbot - Latest Updates:**

* **Multiple Model Support:** You can now choose from different models like Zephyr 7B.
* **Improved Error Handling:** The chatbot now provides clearer error messages if something goes wrong.
"""

def show_updates_and_respond(history, system_message, max_tokens, temperature, top_p, model_name):
    history.append(("User: ", "Show me the latest updates"))
    yield from respond(
        message="Show me the latest updates",
        history=history,
        system_message=system_message,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        model_name=model_name,
    )
    history[-1] = ("User: ", "Show me the latest updates")
    history.append(("Assistant:", latest_updates))
    yield from respond(
        message="What are the latest updates?",
        history=history,
        system_message=system_message,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        model_name=model_name,
    )

# Define the Gradio interface with the Blocks context
with gr.Blocks(css=".gradio-container {border: none;}") as demo:
    chat_history = gr.State([])  # Initialize an empty chat history state
    chat_interface = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Textbox(
                value="You are a friendly and helpful assistant.",
                label="System message",
                lines=2
            ),
            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=2.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)",
            ),
            gr.Dropdown(
                choices=list(available_models.keys()),
                value="Zephyr 7B Beta",
                label="Select Model",
            ),
        ],
        title="Multi-Model Chatbot",
        description="A customizable chatbot interface using Hugging Face's Inference API.",
        chat_history=chat_history,  # Pass the state to the ChatInterface
    )

    # Add the "Show Updates" button and output area
    with gr.Row():
        updates_button = gr.Button("Show Latest Updates")

    # Define the button's click event (now inside the Blocks context)
    updates_button.click(
        fn=show_updates_and_respond,
        inputs=[chat_history, chat_interface.textbox, gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"), gr.Slider(minimum=0.1, maximum=2.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)"), chat_interface.dropdown], 
        outputs=chat_history
    )

# Launch the Gradio interface in full screen
if __name__ == "__main__":
    demo.launch(share=True, fullscreen=True)