import os
import streamlit as st
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image

# =========================
# CONFIG
# =========================
CHECKPOINT_PATH = "best_binary_model_moderation_last.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# =========================
# MODEL
# =========================
def build_model():
    backbone = models.resnet18(weights=None)  # مهم هنا None لأننا هنحمّل weights من checkpoint
    in_features = backbone.fc.in_features
    backbone.fc = nn.Sequential(
        nn.Dropout(p=0.20),
        nn.Linear(in_features, 1)
    )
    return backbone


# =========================
# DECISION LOGIC
# =========================
def moderation_decision(prob: float, review_threshold: float, block_threshold: float) -> str:
    if prob >= block_threshold:
        return "BLOCK"
    if prob >= review_threshold:
        return "REVIEW"
    return "ALLOW"


# =========================
# IMAGE PREP
# =========================
def load_image_for_inference(image: Image.Image, img_size=224):
    tf = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.CenterCrop((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=(0.485, 0.456, 0.406),
            std=(0.229, 0.224, 0.225)
        ),
    ])
    x = tf(image.convert("RGB")).unsqueeze(0)
    return x


# =========================
# LOAD CHECKPOINT
# =========================
@st.cache_resource
def load_model_and_checkpoint():
    if not os.path.exists(CHECKPOINT_PATH):
        raise FileNotFoundError(f"Checkpoint not found: {CHECKPOINT_PATH}")

    checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)

    model = build_model()
    model.load_state_dict(checkpoint["model_state"])
    model = model.to(DEVICE)
    model.eval()

    block_threshold = float(checkpoint.get("block_threshold", 0.5))
    review_threshold = float(checkpoint.get("review_threshold", 0.35))
    img_size = int(checkpoint.get("img_size", 224))

    return model, block_threshold, review_threshold, img_size


# =========================
# PREDICT
# =========================
@torch.no_grad()
def predict_image(image: Image.Image):
    model, block_threshold, review_threshold, img_size = load_model_and_checkpoint()

    x = load_image_for_inference(image, img_size=img_size).to(DEVICE)

    logits = model(x).squeeze(1)
    prob = torch.sigmoid(logits).item()

    decision = moderation_decision(
        prob=prob,
        review_threshold=review_threshold,
        block_threshold=block_threshold
    )

    return {
        "prob_sensitive": prob,
        "decision": decision,
        "block_threshold": block_threshold,
        "review_threshold": review_threshold,
    }


# =========================
# UI
# =========================
st.set_page_config(page_title="Content Moderation UI", page_icon="🛡️", layout="centered")

st.title("🛡️AI-based Content Moderation system")
uploaded_file = st.file_uploader(
    "Upload an image",
    type=["jpg", "jpeg", "png", "webp"]
)

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    if st.button("Predict"):
        with st.spinner("Analyzing image..."):
            result = predict_image(image)

        prob = result["prob_sensitive"]
        decision = result["decision"]
        block_threshold = result["block_threshold"]
        review_threshold = result["review_threshold"]

        st.subheader("Result")

        if decision == "BLOCK":
            st.error(f"Decision: {decision}")
        elif decision == "REVIEW":
            st.warning(f"Decision: {decision}")
        else:
            st.success(f"Decision: {decision}")

        st.write(f"**Sensitive Probability:** `{prob:.6f}`")
        st.write(f"**Review Threshold:** `{review_threshold:.2f}`")
        #st.write(f"**Block Threshold:** `{block_threshold:.2f}`")

        st.progress(min(max(float(prob), 0.0), 1.0))

        st.json({
            "prob_sensitive": round(prob, 6),
            "decision": decision,
            "review_threshold": review_threshold,
#            "block_threshold": block_threshold,
        })