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])

    # Prepare the conversation history for the API call
    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:
        # Generate a response from the model with streaming
        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 (this is just placeholder text)
latest_updates = """
**Chatbot - Latest Updates:**

* Multiple Model Support
* Improved Error Handling
* Enhanced System Message Input
"""

def show_updates_and_respond(history, system_message, max_tokens, temperature, top_p, model_name):
    """
    Shows the latest updates and generates a response from the model based on the updates.
    """
    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
    system_message = gr.Textbox(
        value="You are a friendly and helpful assistant.",
        label="System message",
        lines=2
    )
    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
    model_dropdown = gr.Dropdown(choices=list(available_models.keys()), value="Zephyr 7B Beta", label="Select Model")
    
    # Define the chat interface
    chatbot = gr.Chatbot()

    # Add the "Show Updates" button
    updates_button = gr.Button("Show Latest Updates")

    # Button click event
    updates_button.click(
        fn=show_updates_and_respond,
        inputs=[chat_history, system_message, max_tokens, temperature, top_p, model_dropdown],
        outputs=chat_history
    )

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch(share=True)