import gradio as gr
from openai import OpenAI
import re

def get_openrouter_client(api_key):
    """Initialize OpenRouter client with user-provided API key"""
    if not api_key or api_key.strip() == "":
        return None, "Please enter your OpenRouter API key"
    
    try:
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key
        )
        return client, None
    except Exception as e:
        return None, f"Error initializing client: {str(e)}"

def extract_medicine_names(image, api_key):
    """Extract medicine names from a prescription image using Gemini via OpenRouter"""
    if not image:
        return "Please upload a prescription image."
    
    # Get client with user-provided API key
    client, error = get_openrouter_client(api_key)
    if error:
        return error
    
    try:
        response = client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "https://medicine-extractor-app.com",
                "X-Title": "Medicine Name Extractor",
            },
            model="google/gemini-2.5-pro-exp-03-25:free",
            messages=[
                {
                    "role": "system",
                    "content": "You are an AI specialized in extracting medication names from prescription images. Only list the medication names, nothing else."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Extract ONLY the names of medications from this prescription image. Provide them as a numbered list. If this isn't a medical prescription, respond with 'No prescription detected'."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": image
                            }
                        }
                    ]
                }
            ],
            max_tokens=300
        )
        
        result = response.choices[0].message.content.strip()
        
        # Check if no prescription was detected
        if "No prescription detected" in result:
            return "No prescription detected in the image."
        
        # Clean up the response to just include the medication names
        # Remove any explanatory text that might appear before or after the list
        medicines = []
        for line in result.split('\n'):
            # Look for numbered lines or lines starting with medication names
            if re.match(r'^\d+\.', line.strip()):
                # Extract text after the number and period
                med_name = re.sub(r'^\d+\.\s*', '', line.strip())
                medicines.append(med_name)
        
        if not medicines:
            # If numbered list processing didn't work, return the raw output
            return result
        
        return "\n".join([f"{i+1}. {med}" for i, med in enumerate(medicines)])
        
    except Exception as e:
        return f"Error: {str(e)}"

# Create the Gradio interface
with gr.Blocks(title="Prescription Medicine Extractor") as app:
    gr.Markdown("# Prescription Medicine Name Extractor")
    gr.Markdown("Upload a prescription image to extract medication names.")
    
    api_key = gr.Textbox(
        label="OpenRouter API Key", 
        placeholder="Enter your OpenRouter API key here", 
        type="password"
    )
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Prescription Image")
            submit_btn = gr.Button("Extract Medicine Names", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Extracted Medicine Names", lines=10)
    
    submit_btn.click(
        fn=extract_medicine_names, 
        inputs=[image_input, api_key], 
        outputs=[output]
    )
    
    gr.Markdown("""
    ## Usage Instructions
    1. Enter your OpenRouter API key (get one from https://openrouter.ai)
    2. Upload a clear image of a medical prescription
    3. Click the "Extract Medicine Names" button
    4. The names of medications will be displayed in the output box
    
    **Note:** For best results, ensure the image is clear and the text is readable.
    **Privacy Notice:** Your API key and images are processed only during the active session and are not stored.
    """)

# Launch the app
if __name__ == "__main__":
    print("Starting Prescription Medicine Name Extractor application...")
    app.launch()