import streamlit as st
from transformers import pipeline
from PIL import Image

# Sidebar for module selection
module = st.sidebar.selectbox("Choose Detection Type", ["Text Misinformation", "Image Deepfake"])

# Load models directly from Hugging Face (no local saving)
@st.cache_resource
def load_text_model():
    return pipeline("text-classification", model="mrm8488/bert-tiny-finetuned-fake-news-detection")

@st.cache_resource
def load_image_model():
    return pipeline("image-classification", model="prithivMLmods/Deep-Fake-Detector-v2-Model")

# Text Misinformation Detection
if module == "Text Misinformation":
    st.title("📰 Misinformation Detection")
    user_input = st.text_area("Enter a news statement or claim:")
    if st.button("Analyze"):
        if user_input.strip():
            model = load_text_model()
            result = model(user_input)[0]
            st.success(f"Prediction: {result['label']}")
            st.write(f"Confidence: {result['score']:.2f}")
        else:
            st.warning("Please enter some text.")


# Image Deepfake Detection
elif module == "Image Deepfake":
    st.title("🖼️ Deepfake Image Detection")
    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
    if uploaded_file:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_column_width=True)
        model = load_image_model()
        result = model(image)[0]

        # Map label_0 and label_1 to meaningful labels with icons and colors
        label_map = {
            "LABEL_0": ("Unauthentic", "❌", "red"),
            "LABEL_1": ("Authentic", "✅", "green")
        }
        label, icon, color = label_map.get(result['label'].upper(), (result['label'], "❓", "gray"))

        st.markdown(
            f"<h4 style='color:{color};'>{icon} Prediction: {label}</h4>",
            unsafe_allow_html=True
        )
        st.write(f"Confidence: {result['score']:.2f}")