import gradio as gr
import requests
import base64
import io
import time
import os
from PIL import Image

# ==========================
# CONFIGURATION
# ==========================
MODEL_ID = "gemini-2.5-pro"
API_URL = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL_ID}:generateContent"
MAX_RETRIES = 3
GEMINI_ENV_KEY = "GEMINI_API_KEY"

SYSTEM_INSTRUCTION = (
"""
You are an AI dermatology assistant designed to help patients understand possible skin conditions from images. 
Your goal is to offer an informative, empathetic explanation in clear, simple language. 
Avoid technical jargon. 
Keep your answer brief and concise.
No need to give any disclaimer or introduction.
Generate answers in reportable Bengali language.

"""
)

USER_QUERY = (
"""
Analyze this image of a skin lesion and provide:
1. A brief description of what you observed in the image that supports your conclusion.
2. The diagnosis as Disease Name.
3. A brief explanation of the condition for the patient — including typical symptoms, possible causes, and self-care advice.
"""
)


# ==========================
# HELPER FUNCTIONS
# ==========================
def pil_to_base64(image: Image.Image) -> str:
    """Convert PIL Image to base64-encoded JPEG."""
    buffer = io.BytesIO()
    if image.mode != "RGB":
        image = image.convert("RGB")
    image.save(buffer, format="JPEG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def generate_gemini_analysis(image: Image.Image, max_output_tokens=2048, temperature=0.0):
    """Send image to Gemini API and return dermatological analysis with a live spinner and status updates."""
    def build_spinner_html(status: str = "Working...", hidden: bool = False) -> str:
        display = "none" if hidden else "flex"
        return f"""
        <div style="display:{display};align-items:center;gap:12px;justify-content:center;">
            <div style="display:flex;flex-direction:column;align-items:center;gap:8px;">
                <div style="
                    width:48px;height:48px;
                    border:6px solid rgba(255,255,255,0.12);
                    border-top-color:#06b6d4;
                    border-radius:50%;
                    animation: spin 1s linear infinite;
                "></div>
                <div style="color:#ddd;font-size:0.95em;min-width:180px;text-align:center;">{status}</div>
            </div>
        </div>
        <style>
        @keyframes spin {{ from {{ transform: rotate(0deg); }} to {{ transform: rotate(360deg); }} }}
        </style>
        """

    # initial spinner
    yield build_spinner_html("Initializing..."), ""

    api_key = os.environ.get(GEMINI_ENV_KEY)
    if not api_key:
        # hide spinner when showing error/result
        yield build_spinner_html("Missing API key", hidden=True), "⚠️ **Missing API key.** Set GEMINI_API_KEY as an environment variable."
        return

    if not image:
        yield build_spinner_html("No image provided", hidden=True), "⚠️ Please upload an image."
        return

    yield build_spinner_html("Preparing image..."), ""

    image_b64 = pil_to_base64(image)

    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [
                    {"text": USER_QUERY},
                    {"inlineData": {"mimeType": "image/jpeg", "data": image_b64}},
                ],
            }
        ],
        "systemInstruction": {"parts": [{"text": SYSTEM_INSTRUCTION}]},
        "generationConfig": {"temperature": temperature, "maxOutputTokens": max_output_tokens},
    }

    headers = {"Content-Type": "application/json"}

    # before sending
    yield build_spinner_html("Sending to Gemini..."), ""

    # Retry logic
    for attempt in range(MAX_RETRIES):
        try:
            yield build_spinner_html("Working..."), ""
            response = requests.post(
                f"{API_URL}?key={api_key}",
                headers=headers,
                json=payload,
                timeout=90,
            )
            response.raise_for_status()
            data = response.json()
            break
        except requests.exceptions.RequestException as e:
            if attempt < MAX_RETRIES - 1:
                yield build_spinner_html("Network error — retrying..."), ""
                time.sleep(2 ** attempt)
                continue
            yield build_spinner_html("Request failed", hidden=True), f"❌ Request failed: {e}"
            return

    yield build_spinner_html("Parsing response..."), ""

    candidate = data.get("candidates", [{}])[0]
    parts = candidate.get("content", {}).get("parts", [])
    generated_text = "".join(part.get("text", "") for part in parts)

    if not generated_text:
        finish_reason = candidate.get("finishReason", "")
        yield build_spinner_html("No output", hidden=True), f"⚠️ No text output from model. Finish reason: `{finish_reason}`"
        return

    final_md = f"### 🩺 Dermatological Result \n\n{generated_text.strip()}"

    # hide spinner, show final result
    yield build_spinner_html("Completed", hidden=True), final_md


# ==========================
# CUSTOM STYLING
# ==========================
css = """
#container {
    margin: 0 auto;
    max-width: 600px;
    text-align: center;
}

#analyze-button {
    display: flex;
    justify-content: center;
    margin-top: 15px;
}

#analyze-button button {
    width: auto !important;
    padding: 6px 18px !important;
    border-radius: 10px !important;
    font-weight: 600;
}

.sample-img {
    margin-top: 25px;
    text-align: center;
    opacity: 0.9;
}

.sample-img img {
    border-radius: 10px;
    width: 220px;
    border: 1px solid #444;
}
"""

# ==========================
# GRADIO UI
# ==========================
with gr.Blocks(css=css, title="Skin Disease Analyzer") as demo:
    with gr.Column(elem_id="container"):
        gr.Markdown(
            """
# Skin Disease Test  
**এআই ডাক্তারকে আক্রান্ত স্কিনের একটি ছবি দিন**
            """
        )

        image_input = gr.Image(
            label="Upload Image",
            type="pil",
            show_label=False,
            height=350,
        )

        # Centered analyze button (smaller)
        with gr.Row(elem_id="analyze-button"):
            analyze_btn = gr.Button("Test", variant="primary")

        # Progress HTML + Result Markdown
        progress_html = gr.HTML()
        output_md = gr.Markdown(label="Result")

        # Working sample image (local file)
        gr.Markdown('<p style="font-size:0.9em;margin-top:50px;">📸 Sample Skin Lesion</p>')
        gr.Image(value="assets/bdsd.jpg", show_label=False, interactive=False, height=300)

        # Event binding
        analyze_btn.click(
            fn=generate_gemini_analysis,
            inputs=[image_input],
            outputs=[progress_html, output_md],
        )

demo.launch()