import gradio as gr
from ultralytics import YOLO
from PIL import Image, ImageOps, ImageEnhance
import numpy as np
import tempfile, io, base64

# Load models
model_swelling = YOLO("model/swelling/best.pt")
model_redness = YOLO("model/redness/best.pt")
model_bleeding = YOLO("model/bleeding/best.pt")


# --- Helpers ---
def preprocess(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    image = ImageOps.exif_transpose(image).convert("RGB")

    # Resize if needed
    w, h = image.size
    max_dim = max(w, h)
    if max_dim > 1024:
        scale = 1024 / max_dim
        image = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS)

    # Light contrast boost
    image = ImageEnhance.Contrast(image).enhance(1.05)

    return image


def np_to_base64(img_np, format="JPEG"):
    """Convert numpy RGB image to Base64 string."""
    pil_img = Image.fromarray(img_np)
    buffer = io.BytesIO()
    pil_img.save(buffer, format=format)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def base64_to_pil(b64_str):
    """Convert Base64 string back to PIL image (for Gradio display)."""
    return Image.open(io.BytesIO(base64.b64decode(b64_str)))


# --- Main detection ---
def detect_gingivitis(image, conf=0.4, iou=0.5):
    image = preprocess(image)

    sw_res = model_swelling.predict(image, conf=conf, iou=iou)
    rd_res = model_redness.predict(image, conf=conf, iou=iou)
    bl_res = model_bleeding.predict(image, conf=conf, iou=iou)

    # Convert YOLO output → numpy RGB
    img_sw = sw_res[0].plot()[:, :, ::-1]  # BGR → RGB
    img_rd = rd_res[0].plot()[:, :, ::-1]
    img_bl = bl_res[0].plot()[:, :, ::-1]

    # Encode images to Base64 (for backend API consumption)
    sw_b64 = np_to_base64(img_sw)
    rd_b64 = np_to_base64(img_rd)
    bl_b64 = np_to_base64(img_bl)

    # Convert Base64 back to PIL for Gradio display
    sw_pil = base64_to_pil(sw_b64)
    rd_pil = base64_to_pil(rd_b64)
    bl_pil = base64_to_pil(bl_b64)

    # Determine diagnosis
    has_sw = len(sw_res[0].boxes) > 0
    has_rd = len(rd_res[0].boxes) > 0
    has_bl = len(bl_res[0].boxes) > 0

    if has_sw and has_rd and has_bl:
        diagnosis = (
            "🦷 You have gingivitis.\n\n"
            "For accurate assessment and guidance, we recommend visiting your dentist.\n"
            "If you have a periapical X-ray, you may try the *Detect Periodontitis* tool."
        )
    else:
        diagnosis = "🟢 You don't have gingivitis."

    # Return PIL for Gradio + Base64 is available for backend
    return [sw_pil, rd_pil, bl_pil, diagnosis]


# --- Gradio Interface ---
interface = gr.Interface(
    fn=detect_gingivitis,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(0, 1, value=0.4, step=0.05, label="Confidence Threshold"),
        gr.Slider(0, 1, value=0.5, step=0.05, label="NMS IoU Threshold"),
    ],
    outputs=[
        gr.Image(label="Swelling Detection", type="pil"),
        gr.Image(label="Redness Detection", type="pil"),
        gr.Image(label="Bleeding Detection", type="pil"),
        gr.Textbox(label="Diagnosis")
    ],
    title="Gingivitis Detection"
)

if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860, show_error=True)