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

# =========================================================
# Lazy-loaded global models (LOAD ONLY ON FIRST REQUEST)
# =========================================================
model_swelling = None
model_redness = None
model_bleeding = None

def get_models():
    """Load YOLO models only once (lazy loading)."""
    global model_swelling, model_redness, model_bleeding

    if model_swelling is None:
        model_swelling = YOLO("models/swelling/swelling_final.pt")

    if model_redness is None:
        model_redness = YOLO("models/redness/redness_final.pt")

    if model_bleeding is None:
        model_bleeding = YOLO("models/bleeding/bleeding_final.pt")

    return model_swelling, model_redness, model_bleeding


# =========================================================
# Helper functions
# =========================================================
def preprocess(image):
    """Resize, fix orientation, improve contrast."""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

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

    # Resize if too large
    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)

    # Slight contrast enhancement
    image = ImageEnhance.Contrast(image).enhance(1.05)

    return image


def np_to_base64(img_np, format="JPEG"):
    """Convert numpy RGB image to Base64."""
    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 to PIL image."""
    return Image.open(io.BytesIO(base64.b64decode(b64_str)))


# =========================================================
# Main detection function
# =========================================================
def detect_gingivitis(image, conf=0.25, iou=0.5):
    try:
        if image is None:
            return [None, None, None, "❌ No image uploaded"]

        # Load models (only once)
        sw_model, rd_model, bl_model = get_models()

        # Preprocess
        image = preprocess(image)

        # Run detections
        sw_res = sw_model.predict(image, conf=conf, iou=iou)
        rd_res = rd_model.predict(image, conf=conf, iou=iou)
        bl_res = bl_model.predict(image, conf=conf, iou=iou)

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

        sw_pil = base64_to_pil(np_to_base64(img_sw))
        rd_pil = base64_to_pil(np_to_base64(img_rd))
        bl_pil = base64_to_pil(np_to_base64(img_bl))

        # Diagnosis logic
        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\n"
                "If you have a periapical X-ray, you may try the Detect Periodontitis tool."
            )
        else:
            diagnosis = "🟢 You don't have gingivitis."

        return [sw_pil, rd_pil, bl_pil, diagnosis]

    except Exception as e:
        return [None, None, None, f"❌ Error during processing: {str(e)}"]


# =========================================================
# Gradio Interface
# =========================================================
interface = gr.Interface(
    fn=detect_gingivitis,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(0, 1, value=0.5, 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"
)


# =========================================================
# Warm-start: preload models on startup
# =========================================================
print("🔥 Preloading models to reduce Render cold start...")
get_models()
print("✅ Gingivitis models ready")


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