import gradio as gr
import requests
import json

def get_chat_response(api_key, user_input, history):
    """
    Sends the user input to the Fireworks AI API using the provided API key
    and retrieves the assistant's response.

    Args:
        api_key (str): User-provided Fireworks AI API key.
        user_input (str): The message input by the user.
        history (list): The conversation history.

    Returns:
        tuple: Updated conversation history with the assistant's response.
    """
    url = "https://api.fireworks.ai/inference/v1/chat/completions"

    payload = {
        "model": "accounts/fireworks/models/llama-v3p1-405b-instruct",
        "max_tokens": 16384,
        "top_p": 1,
        "top_k": 40,
        "presence_penalty": 0,
        "frequency_penalty": 0,
        "temperature": 0.6,
        "messages": [
            {
                "role": "user",
                "content": user_input
            }
        ]
    }

    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload))
        response.raise_for_status()  # Raise an exception for HTTP errors
        response_data = response.json()

        # Extract the assistant's reply
        if "choices" in response_data and len(response_data["choices"]) > 0:
            assistant_reply = response_data["choices"][0]["message"]["content"]
        else:
            assistant_reply = "No response from the assistant."

    except requests.exceptions.HTTPError as http_err:
        assistant_reply = f"HTTP error occurred: {http_err}"
    except Exception as err:
        assistant_reply = f"An error occurred: {err}"

    # Append the user input and assistant reply to the history
    history = history + [[user_input, assistant_reply]]
    return history, history

def set_api_key(api_key):
    """
    Stores the user's API key.

    Args:
        api_key (str): User-provided Fireworks AI API key.

    Returns:
        tuple: Stored API key and confirmation message.
    """
    if api_key:
        return api_key, "✅ API Key has been set."
    else:
        return "", "❌ Please enter a valid API key."

def send_message(user_input, history, api_key):
    """
    Handles sending messages to the chatbot.

    Args:
        user_input (str): The message input by the user.
        history (list): The conversation history.
        api_key (str): User-provided Fireworks AI API key.

    Returns:
        tuple: Updated conversation history with the assistant's response.
    """
    if not api_key:
        assistant_reply = "❌ Please enter your Fireworks AI API key above."
        history = history + [[user_input, assistant_reply]]
        return history, history

    return get_chat_response(api_key, user_input, history)

# Define the Gradio Blocks interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🪄 Fireworks AI Chatbot
        Interact with the Fireworks AI chatbot powered by the LLaMA model.
        Enter your **Fireworks AI API key** and your messages below to receive responses in real-time.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=7):
            api_key_input = gr.Textbox(
                label="🔑 Fireworks AI API Key",
                placeholder="Enter your API key here...",
                type="password",
                lines=1
            )
        with gr.Column(scale=3):
            set_key_button = gr.Button("Set API Key")
    
    # Output for API key confirmation
    api_key_status = gr.Textbox(
        label="API Key Status",
        interactive=False,
        lines=1,
        value=""
    )
    
    chatbot = gr.Chatbot(label="💬 Chatbot")
    
    with gr.Row():
        with gr.Column(scale=8):
            user_input = gr.Textbox(
                placeholder="Type your message here...",
                label="📝 Your Message",
                lines=1
            )
        with gr.Column(scale=2):
            send_button = gr.Button("Send")
    
    # Optional: Add example buttons
    with gr.Accordion("💡 Examples", open=False):
        examples = gr.Examples(
            examples=[
                "Tell me a joke.",
                "Explain the theory of relativity.",
                "How does photosynthesis work?",
                "What's the capital of France?"
            ],
            inputs=user_input,
            label="Click an example to use it"
        )
    
    # Initialize the conversation history and store the API key
    history = gr.State([])
    stored_api_key = gr.State("")
    
    # Set API Key Button Click Event
    set_key_button.click(
        fn=set_api_key,
        inputs=api_key_input,
        outputs=[stored_api_key, api_key_status]
    )
    
    # Define the function to handle the interaction
    send_button.click(
        fn=send_message,
        inputs=[user_input, history, stored_api_key],
        outputs=[chatbot, history]
    )
    
    # Allow pressing Enter to send the message
    user_input.submit(
        fn=send_message,
        inputs=[user_input, history, stored_api_key],
        outputs=[chatbot, history]
    )
    
    gr.Markdown(
        """
        ---
        **ℹ️ Note**: Ensure you have a valid Fireworks AI API key. You can obtain one from [Fireworks AI](https://fireworks.ai/).
        """
    )

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