import streamlit as st
import requests
from PIL import Image

# ==============================
# CONFIG
# ==============================
API_BASE = "https://AdarshDS-mold-detection-api.hf.space"
CONVNEXT_API = f"{API_BASE}/predict/v2"

st.set_page_config(
    page_title="Mold Detection System",
    layout="centered"
)

# ==============================
# UI HEADER
# ==============================
st.title("🦠 AI Mold Detection System")
st.caption("Powered by ConvNeXt (Advanced Model)")

st.write(
    "Upload an image of a wall or ceiling. "
    "The image is analyzed using an advanced ConvNeXt-based AI model "
    "with uncertainty estimation and self-supervised verification."
)

file = st.file_uploader(
    "Upload Image",
    type=["jpg", "png", "jpeg"]
)

# ==============================
# MAIN LOGIC
# ==============================
if file:
    image = Image.open(file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    if st.button("Analyze"):
        file_bytes = file.getvalue()
        files_payload = {
            "file": (file.name, file_bytes, file.type)
        }

        with st.spinner("🔍 Analyzing image..."):
            resp = requests.post(CONVNEXT_API, files=files_payload)

        st.markdown("---")
        st.subheader("📊 Prediction Result")

        if resp.status_code != 200:
            st.error("❌ API error while analyzing the image.")
            st.text(resp.text)
        else:
            res = resp.json()

            # ==============================
            # Main Results
            # ==============================
            st.metric("Decision", res["decision"])
            st.metric(
                "Mold Probability",
                res["model_outputs"]["mold_probability"]
            )
            st.metric(
                "Biological Probability",
                res["model_outputs"]["biological_probability"]
            )

            # ==============================
            # Confidence Checks
            # ==============================
            st.subheader("📈 Confidence Checks")
            cc = res["confidence_checks"]

            c1, c2, c3 = st.columns(3)
            c1.metric("Uncertainty", cc["uncertainty"])
            c2.metric("Patch Ratio", cc["patch_ratio"])
            c3.metric("DINO Similarity", cc["dino_similarity"])

            # ==============================
            # User Feedback
            # ==============================
            if res["decision"] == "Mold":
                st.error(
                    "❌ Mold detected. "
                    "Professional remediation is strongly recommended."
                )
            elif res["decision"] == "Possible Mold":
                st.warning(
                    "⚠️ Possible mold detected. "
                    "Human inspection is advised."
                )
            else:
                st.success("✅ No mold detected.")