import gradio as gr
from huggingface_hub import InferenceClient

# Define the InferenceClient for different models
client_chatgpt = InferenceClient("openai/gpt-3.5-turbo")  # Example for ChatGPT model
client_llama = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")  # Llama model
client_claude = InferenceClient("anthropic/claude-1")  # Claude model (adjust with correct model path)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    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})

    response = ""

    for message in client_llama.chat_completion(  # Defaulting to Llama, update dynamically later
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

# Function to handle button clicks for different models
def on_button_click(model_name, message, history, system_message, max_tokens, temperature, top_p):
    # Choose the client based on the selected model
    if model_name == "Chatgpt":
        client = client_chatgpt
    elif model_name == "Llama":
        client = client_llama
    elif model_name == "Claude":
        client = client_claude
    else:
        return "Unknown model selected."

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

    response = ""

    # Call the selected model for completion
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# CSS for styling the interface
css = """
body {
    background-color: #06688E; /* Dark background */
    color: white; /* Text color for better visibility */
}
.gr-button {
    background-color: #42B3CE !important; /* White button color */
    color: black !important; /* Black text for contrast */
    border: none !important;
    padding: 8px 16px !important;
    border-radius: 5px !important;
}
.gr-button:hover {
    background-color: #e0e0e0 !important; /* Slightly lighter button on hover */
}
.gr-slider-container {
    color: white !important; /* Slider labels in white */
}
"""

# Interface using Blocks context
with gr.Blocks() as demo:
    # Add all your components here, including buttons
    system_message = gr.Textbox(value="You are a virtual health assistant...", label="System message", visible=False)
    message_input = gr.Textbox(label="User message")
    history = gr.State([])  # Keep the history of interactions
    max_tokens_slider = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
    temperature_slider = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
    top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")

    # Buttons to select the model
    gr.Button("Chatgpt").click(on_button_click, inputs=[message_input, history, system_message, max_tokens_slider, temperature_slider, top_p_slider], outputs="text")
    gr.Button("Llama").click(on_button_click, inputs=[message_input, history, system_message, max_tokens_slider, temperature_slider, top_p_slider], outputs="text")
    gr.Button("Claude").click(on_button_click, inputs=[message_input, history, system_message, max_tokens_slider, temperature_slider, top_p_slider], outputs="text")
    
    # Optional: customize your layout with CSS if needed
    demo.css = css

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