import os
import gradio as gr
from huggingface_hub import login
from huggingface_hub import InferenceClient
import spaces

# Authenticate with Hugging Face API
api_key = os.getenv("LLAMA")
login(api_key)

# Initialize clients for different models
llama_client = InferenceClient("meta-llama/Llama-3.1-70B-Instruct")
gpt_client = InferenceClient("openai/gpt-4")  # Example: Replace with your OpenAI GPT model

# Define the response function
@spaces.GPU
def respond(
    message,
    history: list[dict],
    system_message,
    max_tokens,
    temperature,
    top_p,
    selected_models,
):
    # Prepare input messages
    messages = [{"role": "system", "content": system_message}] + history
    messages.append({"role": "user", "content": message})

    # Collect responses from selected models
    responses = {}

    if "Llama" in selected_models:
        llama_response = ""
        for token in llama_client.chat_completion(
            messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
        ):
            delta = token.choices[0].delta.content
            llama_response += delta
        responses["Llama"] = llama_response

    if "GPT" in selected_models:
        gpt_response = ""
        for token in gpt_client.chat_completion(
            messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
        ):
            delta = token.choices[0].delta.content
            gpt_response += delta
        responses["GPT"] = gpt_response

    return responses

# Build the Gradio app
def create_demo():
    return gr.Blocks().add(
        gr.Markdown("# AI Model Comparison Tool 🌟"),
        gr.ChatInterface(
            respond,
            type="messages",
            additional_inputs=[
                gr.Textbox(
                    value="You are a helpful assistant providing answers for technical and customer support queries.",
                    label="System message"
                ),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max 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)"
                ),
                gr.CheckboxGroup(
                    ["Llama", "GPT"],
                    label="Select models to compare",
                    value=["Llama"]
                ),
            ],
        ),
    )

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