import gradio as gr
import base64
from groq import Groq
import tempfile
import os

def encode_image(image_path):
    """Convert an image to base64 encoding"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def extract_medicines(image, api_key):
    """Extract medicine names from prescription image using Groq API"""
    if not api_key:
        return "Please provide a Groq API key"
    
    if image is None:
        return "Please upload a prescription image"
    
    try:
        # Save the uploaded image to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_image:
            image_path = temp_image.name
            image.save(image_path)
        
        # Encode the image to base64
        base64_image = encode_image(image_path)
        
        # Clean up the temporary file
        os.unlink(image_path)
        
        # Initialize Groq client with the provided API key
        client = Groq(api_key=api_key)
        
        # Create the prompt specifically asking for medicine names only
        prompt = "This is an image of a medical prescription. Extract and list ONLY the names of medicines/drugs/medications from this prescription. Do not include dosages, frequencies, or other information. Return just a simple list of medicine names, one per line."
        
        # Make the API call
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}",
                            },
                        },
                    ],
                }
            ],
            model="meta-llama/llama-4-scout-17b-16e-instruct",
        )
        
        # Return the extracted medicine names
        return chat_completion.choices[0].message.content
    
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
with gr.Blocks(title="Prescription Medicine Extractor") as app:
    gr.Markdown("# Medicine Name Extractor from Prescriptions")
    gr.Markdown("Upload a prescription image and enter your Groq API key to extract medicine names")
    
    with gr.Row():
        with gr.Column():
            api_key_input = gr.Textbox(
                label="Groq API Key", 
                placeholder="Enter your Groq API key here", 
                type="password"
            )
            image_input = gr.Image(label="Upload Prescription Image", type="pil")
            extract_button = gr.Button("Extract Medicine Names")
        
        with gr.Column():
            output = gr.Textbox(label="Extracted Medicine Names", lines=10)
    
    extract_button.click(
        fn=extract_medicines, 
        inputs=[image_input, api_key_input], 
        outputs=output
    )
    
    gr.Markdown("""
    ## Instructions
    1. Enter your Groq API key
    2. Upload a clear image of a medical prescription
    3. Click "Extract Medicine Names"
    4. The application will return only the names of medicines from the prescription
    
    **Note:** Your API key is not stored and is only used for making requests to the Groq API.
    """)

if __name__ == "__main__":
    app.launch()