import gradio as gr
import requests
import os
from IPython.display import display
from IPython.display import Markdown

api_key = os.environ.get("HYPERBOLIC_API_KEY")
url = "https://api.hyperbolic.xyz/v1/chat/completions"

def hyperbolic(prompt):

    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        data = {
            "messages": [
                {
                    "role": "user",
                    "content": f"{prompt}"
                }
            ],
            "model": "Qwen/Qwen2.5-72B-Instruct",
            "max_tokens": 32768,
            "temperature": 0.2,
            "top_p": 0.9
        }

        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
        json_response = response.json()

        # Extract the response text from the JSON
        if 'choices' in json_response and len(json_response['choices']) > 0:
            return json_response['choices'][0]['message']['content']
        else:
            return f"Error: Unexpected response format from Hyperbolic API: {json_response}"

    except requests.exceptions.RequestException as e:
        return f"Error: Could not connect to Hyperbolic API: {e}"
    except (KeyError, TypeError) as e:
        return f"Error: Problem parsing Hyperbolic API response: {e}"
    except Exception as e:
        return f"An unexpected error occurred: {e}"


def chat(message, history):

    bot_message = hyperbolic(message)
    history.append((message, bot_message))
    display(Markdown(bot_message))
    #return bot_message
    return history

if __name__ == '__main__':
    if not api_key:
        print("Error: HYPERBOLIC_API_KEY environment variable not set.  Please set it before running.")
    else:
        demo = gr.ChatInterface(
            fn=chat,
            title="Hyperbolic Chatbot",
            description="A chatbot powered by the Hyperbolic API.",
            chatbot=gr.Chatbot(height=400)  # Adjust height as needed
        )
        demo.launch()