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

def extract_medicine_names(api_key, image_path, model_choice):
    """
    Extract medicine names from a prescription image using OpenRouter API
    """
    if not api_key or api_key.strip() == "":
        return "Please provide a valid OpenRouter API key."
    
    if not image_path:
        return "Please upload an image."
    
    try:
        # Initialize the OpenAI client with OpenRouter base URL
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
        
        # Select the model based on user's choice
        if model_choice == "Llama 4 Maverick":
            model = "meta-llama/llama-4-maverick:free"
        else:  # Default to Kimi VL
            model = "moonshotai/kimi-vl-a3b-thinking:free"
        
        # Prepare image for API
        with open(image_path, "rb") as image_file:
            import base64
            image_data = base64.b64encode(image_file.read()).decode("utf-8")
            
        # Create the completion
        completion = client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "https://medicine-extractor-app.com",
                "X-Title": "Medicine Name Extractor",
            },
            model=model,
            messages=[
                {
                    "role": "system",
                    "content": "You are a specialized medical assistant. Your task is to analyze prescription images and extract ONLY the names of medicines/medications. Return them as a clear, numbered list without any other commentary. If you cannot identify any medicine names, state that clearly."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Please extract and list ONLY the names of medicines or medications from this prescription image."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_data}"
                            }
                        }
                    ]
                }
            ]
        )
        
        # Get the response
        result = completion.choices[0].message.content
        
        # Process to ensure we're only returning medicine names
        # This is a basic processing step - the system prompt should already help focus the response
        if "no medicine" in result.lower() or "cannot identify" in result.lower():
            return "No medicine names were identified in the prescription image."
        
        return result
        
    except Exception as e:
        return f"Error: {str(e)}"

# Define the Gradio interface
with gr.Blocks(title="Medicine Name Extractor", theme=gr.themes.Soft()) as app:
    gr.Markdown("# Medicine Name Extractor")
    gr.Markdown("Upload a prescription image and the app will extract medication names using AI vision models.")
    
    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(
                label="OpenRouter API Key", 
                placeholder="Enter your OpenRouter API key here",
                type="password"
            )
            
            model_choice = gr.Radio(
                ["Llama 4 Maverick", "Kimi VL"], 
                label="Select AI Model",
                value="Kimi VL"
            )
            
            image_input = gr.Image(
                label="Upload Prescription Image",
                type="filepath"
            )
            
            submit_btn = gr.Button("Extract Medicine Names", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Extracted Medicine Names",
                lines=10
            )
    
    # Handle the submission
    submit_btn.click(
        fn=extract_medicine_names,
        inputs=[api_key, image_input, model_choice],
        outputs=output
    )
    
    gr.Markdown("""
    ## How to use
    1. Enter your OpenRouter API key
    2. Select an AI vision model
    3. Upload a clear image of a medical prescription
    4. Click "Extract Medicine Names"
    
    ## Privacy Note
    Your prescription images are processed securely. No data is stored on our servers.
    """)

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