import os
import base64
import gradio as gr
from openai import OpenAI

# Initialize the Nebius OpenAI client
client = OpenAI(
    base_url="https://api.studio.nebius.ai/v1/",
    api_key=os.environ.get("NEBIUS_API_KEY"),  # Replace with actual API key if not using environment variables
)

# Function to encode the image in Base64
def encode_image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# Function to interact with the Nebius OpenAI API
def analyze_image(image):
    try:
        # Convert the image to Base64
        image_base64 = encode_image_to_base64(image)

        # API request
        response = client.chat.completions.create(
            model="Qwen/Qwen2-VL-72B-Instruct",  # Ensure this model name is correct
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "What’s in this image?"},
                        {
                            "type": "image_base64",
                            "image_base64": image_base64,
                        },
                    ],
                }
            ],
            max_tokens=300,
        )

        # Extract the AI's response
        if response.choices and "message" in response.choices[0]:
            return response.choices[0]["message"]["content"]
        else:
            return "No valid response received from the API."

    except Exception as e:
        return f"Error: {str(e)}"


# Gradio interface for uploading an image
with gr.Blocks() as app:
    gr.Markdown("# Image Analysis with Nebius OpenAI")
    with gr.Row():
        image_input = gr.Image(type="filepath", label="Upload an Image")
        output_text = gr.Textbox(label="AI Response")
    
    analyze_button = gr.Button("Analyze Image")
    analyze_button.click(analyze_image, inputs=image_input, outputs=output_text)

# Launch the Gradio app
if __name__ == "__main__":
    app.launch()