import gradio as gr
import openai
import base64
from PIL import Image
import os

# 🔑 Get OpenAI key from Hugging Face Secrets
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError("Please set OPENAI_API_KEY in Hugging Face Secrets")

openai.api_key = OPENAI_API_KEY

def encode_image(image_path):
    """Convert uploaded image to base64 string"""
    with open(image_path, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode("utf-8")

def classify_medicine(image_path):
    try:
        image_b64 = encode_image(image_path)

        response = openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are an assistant that classifies medicine packaging."},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Classify the uploaded medicine image as 'Generic' or 'Ethical (Branded)'. Only respond with one of these two labels."},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}}
                    ]
                }
            ],
        )

        result = response.choices[0].message.content.strip()
        return image_path, f"✅ Prediction: {result}"

    except Exception as e:
        return None, f"⚠️ Error: {str(e)}"

# 🎨 Custom UI
custom_css = """
.gradio-container {
    background-color: #f9fdfb !important;
    font-family: 'Segoe UI', sans-serif;
    color: #003333 !important;
    padding: 20px;
}
h1 {
    color: #006666 !important;
    text-align: center;
    font-weight: 700;
    margin-bottom: 10px;
}
p {
    text-align: center;
    color: #004d4d !important;
    font-size: 16px;
}
.card {
    background: white;
    border-radius: 16px;
    padding: 20px;
    box-shadow: 0px 4px 15px rgba(0,0,0,0.08);
}
button {
    background-color: #f9fdfb !important;
    color: black !important;
    border-radius: 8px !important;
    font-weight: bold;
    font-size: 15px;
    padding: 10px 20px;
}
button:hover {
    background-color: #f9fdfb !important;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("<h1>💊 Pharma Medicine Classifier</h1><p>Upload a medicine image to check if it is <b>Generic</b> or <b>Ethical (Branded)</b>.</p>")

    with gr.Row(elem_classes="card"):
        with gr.Column():
            input_img = gr.Image(type="filepath", label="📷 Upload Medicine Image", height=250)
            classify_btn = gr.Button("🔍 Classify Medicine", variant="primary")
        with gr.Column():
            preview_img = gr.Image(label="Uploaded Image", interactive=False, height=250)
            output_label = gr.Textbox(label="Result", interactive=False)

    classify_btn.click(fn=classify_medicine, inputs=input_img, outputs=[preview_img, output_label])

demo.launch()