import streamlit as st
import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig, CLIPProcessor, CLIPModel
from peft import PeftModel
from PIL import Image

st.set_page_config(page_title="Multimodal Risk Engine", page_icon="🛡️", layout="wide")

# --- LOAD MODELS (Smart Caching) ---
@st.cache_resource
def load_models():
    print("🔄 Loading Models...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # A. Stage 1 (CLIP)
    clip_id = "openai/clip-vit-base-patch32"
    clip_model = CLIPModel.from_pretrained(clip_id).to(device)
    clip_processor = CLIPProcessor.from_pretrained(clip_id)
    
    # B. Stage 2 (LLaVA)
    model_id = "llava-hf/llava-1.5-7b-hf"
    
    # CPU vs GPU Loading Logic
    if device == "cuda":
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_use_double_quant=True, 
            bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16
        )
        base_model = LlavaForConditionalGeneration.from_pretrained(
            model_id, quantization_config=bnb_config, torch_dtype=torch.float16, device_map="auto"
        )
    else:
        # Fallback for Free CPU Tier (Might be slow/crash but allows build verification)
        base_model = LlavaForConditionalGeneration.from_pretrained(model_id)
    
    adapter_id = "oke39/llava-v1.5-7b-hateful-memes-lora"
    model = PeftModel.from_pretrained(base_model, adapter_id)
    llava_processor = AutoProcessor.from_pretrained(model_id)
    
    return clip_model, clip_processor, model, llava_processor

try:
    clip_model, clip_processor, llava_model, llava_processor = load_models()
    st.toast("✅ System Ready", icon="🚀")
except Exception as e:
    st.error(f"Hardware Error: {e}")
    st.stop()

# --- PIPELINE ---
def stage_1_glance(image):
    HATE_ANCHOR = ["hate speech, offensive content, racism, dangerous propaganda"]
    inputs = clip_processor(text=HATE_ANCHOR, images=image, return_tensors="pt", padding=True).to(clip_model.device)
    with torch.no_grad():
        outputs = clip_model(**inputs)
        probs = outputs.logits_per_image.softmax(dim=1)
        return float(probs[0][0])

def stage_2_judge(image, text_context):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    prompt = f"USER: <image>\nAnalyze this meme text: '{text_context}'. Is it hateful? Return JSON.\nASSISTANT:"
    inputs = llava_processor(text=prompt, images=image, return_tensors="pt").to(device)
    
    with torch.inference_mode():
        output = llava_model.generate(**inputs, max_new_tokens=200)
    response = llava_processor.decode(output[0], skip_special_tokens=True)
    return response.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in response else response

# --- UI ---
st.title("🛡️ Multimodal Content Risk Engine")
st.markdown("### The 'Benign Confounder' Solver")

col1, col2 = st.columns([1, 1])

with col1:
    uploaded_file = st.file_uploader("Upload Meme", type=["jpg", "png", "jpeg"])
    text_input = st.text_input("Extracted Text", placeholder="Type visible text...")
    if uploaded_file:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="User Upload", use_container_width=True)

with col2:
    if uploaded_file and st.button("Analyze Risk", type="primary"):
        with st.status("Running Pipeline...", expanded=True) as status:
            st.write("👁️ **Stage 1: The 'Glance' (CLIP)**")
            risk_score = stage_1_glance(image)
            st.progress(min(risk_score, 1.0))
            
            if risk_score < 0.22:
                status.update(label="✅ Auto-Approved", state="complete")
                st.success("Verdict: BENIGN")
                st.write(f"Risk Score: `{risk_score:.4f}`")
            else:
                st.warning(f"⚠️ Risk Detected ({risk_score:.4f})! Escalating...")
                st.write("⚖️ **Stage 2: The 'Judge' (LLaVA)**")
                verdict = stage_2_judge(image, text_input if text_input else "")
                status.update(label="✅ Done", state="complete")
                st.json(verdict)