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

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ.get("OPENROUTER_API_KEY")
)

def image_to_base64(image):
    pil_image = Image.fromarray(image)
    buffer = io.BytesIO()
    pil_image.save(buffer, format="PNG")
    encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return encoded

def identify_pill(image):
    if image is None:
        return "⚠️ Please upload a medicine image first."

    prompt = """
    You are a pharmaceutical assistant. Analyze this medicine image and provide:

    💊 MEDICINE IDENTIFICATION
    - Medicine name (brand + generic if visible)
    - Manufacturer (if visible)
    - Type (tablet / capsule / syrup / strip)

    🏥 WHAT IT TREATS
    - Primary use / condition it treats
    - Common diseases it is prescribed for

    📋 DOSAGE INFORMATION
    - Standard adult dosage
    - How to take it (with food / without food / with water)
    - Frequency (once/twice/thrice daily)

    ⚠️ SIDE EFFECTS
    - Common side effects
    - Serious side effects to watch for

    🚫 WARNINGS
    - Who should NOT take this medicine
    - Drug interactions to be aware of
    - Pregnancy / children warnings if any

    👨‍⚕️ IMPORTANT
    - Always end with: "Please consult a doctor or pharmacist before taking any medication."

    If you cannot identify the medicine clearly, say so honestly and ask for a clearer image.
    """

    try:
        img_base64 = image_to_base64(image)
        response = client.chat.completions.create(
            model="nex-agi/nex-n2-pro:free",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{img_base64}"
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt
                        }
                    ]
                }
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"❌ Error: {str(e)}"


custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

* { font-family: 'Inter', sans-serif; }

body { background: linear-gradient(135deg, #f0f4ff 0%, #e8f5e9 100%); }

.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}

/* Hero Header */
.hero-header {
    background: linear-gradient(135deg, #1a73e8 0%, #0d47a1 100%);
    border-radius: 20px;
    padding: 40px;
    text-align: center;
    margin-bottom: 24px;
    box-shadow: 0 8px 32px rgba(26,115,232,0.3);
}

.hero-header h1 {
    color: white !important;
    font-size: 2.8em !important;
    font-weight: 700 !important;
    margin: 0 !important;
    letter-spacing: -0.5px;
}

.hero-header p {
    color: rgba(255,255,255,0.85) !important;
    font-size: 1.1em !important;
    margin-top: 8px !important;
}

.badge {
    display: inline-block;
    background: rgba(255,255,255,0.2);
    color: white;
    padding: 4px 14px;
    border-radius: 20px;
    font-size: 0.85em;
    margin-top: 12px;
    border: 1px solid rgba(255,255,255,0.3);
}

/* Upload Card */
.upload-card {
    background: white;
    border-radius: 16px;
    padding: 24px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08);
    border: 1px solid #e3f2fd;
}

/* Result Card */
.result-card {
    background: white;
    border-radius: 16px;
    padding: 24px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08);
    border: 1px solid #e8f5e9;
}

/* Tips Box */
.tips-box {
    background: linear-gradient(135deg, #e3f2fd, #f3e5f5);
    border-radius: 12px;
    padding: 16px 20px;
    margin-top: 16px;
    border-left: 4px solid #1a73e8;
}

.tips-box p {
    margin: 4px 0 !important;
    color: #37474f !important;
    font-size: 0.9em !important;
}

/* Identify Button */
.identify-btn {
    background: linear-gradient(135deg, #1a73e8, #0d47a1) !important;
    border: none !important;
    border-radius: 12px !important;
    font-size: 1.1em !important;
    font-weight: 600 !important;
    padding: 14px !important;
    letter-spacing: 0.3px !important;
    box-shadow: 0 4px 15px rgba(26,115,232,0.4) !important;
    transition: all 0.3s ease !important;
}

.identify-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(26,115,232,0.5) !important;
}

/* Textbox styling */
textarea {
    border-radius: 12px !important;
    border: 1.5px solid #e3f2fd !important;
    font-size: 0.95em !important;
    line-height: 1.7 !important;
    color: #263238 !important;
    background: #fafffe !important;
}

/* Stats row */
.stats-row {
    display: flex;
    gap: 12px;
    margin-bottom: 20px;
}

.stat-card {
    flex: 1;
    background: white;
    border-radius: 12px;
    padding: 16px;
    text-align: center;
    box-shadow: 0 2px 12px rgba(0,0,0,0.06);
    border-top: 3px solid #1a73e8;
}

.stat-card .number {
    font-size: 1.8em;
    font-weight: 700;
    color: #1a73e8;
}

.stat-card .label {
    font-size: 0.8em;
    color: #78909c;
    margin-top: 4px;
}

/* Warning banner */
.warning-banner {
    background: linear-gradient(135deg, #fff3e0, #fce4ec);
    border-radius: 12px;
    padding: 14px 20px;
    margin-top: 20px;
    border-left: 4px solid #ff6f00;
    text-align: center;
}

.warning-banner p {
    color: #bf360c !important;
    font-size: 0.88em !important;
    margin: 0 !important;
    font-weight: 500 !important;
}

/* Image upload area */
.image-upload {
    border-radius: 12px !important;
    border: 2px dashed #90caf9 !important;
}
"""

with gr.Blocks(css=custom_css, title="💊 PillIdentifier") as app:

    # Hero Header
    gr.HTML("""
    <div class="hero-header">
        <h1>💊 PillIdentifier</h1>
        <p>Upload any medicine photo and get instant, detailed drug information</p>
        <span class="badge">🤖 Powered by AI Vision · Free · Instant</span>
    </div>
    """)

    # Stats Row
    gr.HTML("""
    <div class="stats-row">
        <div class="stat-card">
            <div class="number">💊</div>
            <div class="label">Tablets & Capsules</div>
        </div>
        <div class="stat-card">
            <div class="number">📋</div>
            <div class="label">Dosage Info</div>
        </div>
        <div class="stat-card">
            <div class="number">⚠️</div>
            <div class="label">Side Effects</div>
        </div>
        <div class="stat-card">
            <div class="number">🚫</div>
            <div class="label">Warnings</div>
        </div>
    </div>
    """)

    with gr.Row(equal_height=True):

        # Left Column — Upload
        with gr.Column(scale=1):
            gr.HTML('<div class="upload-card">')
            gr.HTML('<h3 style="color:#1a73e8; margin:0 0 16px 0;">📤 Upload Medicine</h3>')
            image_input = gr.Image(
                label="",
                type="numpy",
                height=300,
                elem_classes=["image-upload"]
            )
            identify_btn = gr.Button(
                "🔍 Identify Medicine",
                variant="primary",
                size="lg",
                elem_classes=["identify-btn"]
            )
            gr.HTML("""
            <div class="tips-box">
                <p>📸 <strong>Tips for best results:</strong></p>
                <p>✅ Clear, well-lit photo</p>
                <p>✅ Include medicine name/strip</p>
                <p>✅ Avoid blurry or dark images</p>
                <p>✅ Works with tablets, capsules, strips</p>
            </div>
            """)
            gr.HTML('</div>')

        # Right Column — Result
        with gr.Column(scale=1):
            gr.HTML('<div class="result-card">')
            gr.HTML('<h3 style="color:#2e7d32; margin:0 0 16px 0;">📋 Medicine Information</h3>')
            output = gr.Textbox(
                label="",
                lines=20,
                placeholder="Your medicine analysis will appear here...\n\nUpload a medicine photo and click 'Identify Medicine' to get started!",
            )
            gr.HTML('</div>')

    # Warning Banner
    gr.HTML("""
    <div class="warning-banner">
        <p>⚠️ <strong>Medical Disclaimer:</strong> This app is for informational purposes only.
        Always consult a qualified doctor or pharmacist before taking any medication.
        Do not self-medicate based on AI output alone.</p>
    </div>
    """)

    identify_btn.click(
        fn=identify_pill,
        inputs=[image_input],
        outputs=output
    )

app.launch()