import streamlit as st
import numpy as np
import requests
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ---------------------------
# PAGE CONFIG
# ---------------------------
st.set_page_config(
    page_title="TruthGuard AI",
    layout="wide",
    page_icon="🛡️"
)

st.markdown("""
<style>
.main { background-color: #0E1117; color: white; }
.stButton>button {
    background-color: #4CAF50;
    color: white;
    border-radius: 10px;
    height: 3em;
    width: 100%;
}
</style>
""", unsafe_allow_html=True)

st.title("🛡️ TruthGuard AI")
st.caption("Multi-Modal Fake News & AI Image Detection System")

# ---------------------------
# LOAD TEXT MODEL
# ---------------------------
@st.cache_resource
def load_text_model():
    # model_name = "Maheentouqeer1/truthguard-fake-news-detector"
    # NEW - replace with this well-trained model
    model_name = "hamzab/roberta-fake-news-classification"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name, low_cpu_mem_usage=True
    )
    model.eval()
    return tokenizer, model

# ---------------------------
# LOAD IMAGE MODEL (PyTorch only, no TensorFlow)
# ---------------------------
# @st.cache_resource
# def load_image_model():
#     import os
#     import tensorflow as tf
#     model_path = "image_model.h5"

#     if not os.path.exists(model_path):
#         url = "https://huggingface.co/syeda-Rija20/image-detector/resolve/main/image_detector_finetuned.h5"
#         response = requests.get(url)
#         with open(model_path, "wb") as f:
#             f.write(response.content)

#     model = tf.keras.models.load_model(model_path)
#     return model
@st.cache_resource
def load_image_model():
    import os
    import tensorflow as tf
    import requests

    # Use .keras extension so TF knows it's zip format
    model_path = "/tmp/image_detector_v2.keras"

    if not os.path.exists(model_path):
        url = "https://huggingface.co/syeda-Rija20/image-detector/resolve/main/image_detector_v2.keras?download=true"
        with requests.get(url, stream=True, allow_redirects=True) as r:
            r.raise_for_status()
            with open(model_path, "wb") as f:  # save as .keras not .h5
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)

    model = tf.keras.models.load_model(model_path)
    return model

def predict_image(img, image_model):
    import numpy as np
    import tensorflow as tf
    img_resized = img.resize((224, 224))
    img_array = tf.keras.preprocessing.image.img_to_array(img_resized) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    prediction = image_model.predict(img_array)
    confidence = float(prediction[0][0]) * 100
    label = "AI" if confidence > 50 else "REAL"
    score = confidence if confidence > 50 else 100 - confidence
    return label, score
    

# # ---------------------------
# # PREDICT TEXT
# # ---------------------------
# def predict_news(text, tokenizer, text_model):
#     inputs = tokenizer(
#         text, return_tensors="pt",
#         truncation=True, padding=True, max_length=512
#     )
#     with torch.no_grad():
#         outputs = text_model(**inputs)
#     probs = torch.nn.functional.softmax(outputs.logits, dim=1)
#      # ADD THIS to see raw scores
#     st.write(f"Debug → Label 0 score: {probs[0][0]:.2f}, Label 1 score: {probs[0][1]:.2f}")
#     prediction = torch.argmax(probs).item()
#     confidence = torch.max(probs).item() * 100
#     return prediction, confidence
def predict_news(text, tokenizer, text_model):
    inputs = tokenizer(
        text, return_tensors="pt",
        truncation=True, padding=True, max_length=512
    )
    with torch.no_grad():
        outputs = text_model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
    prediction = torch.argmax(probs).item()
    confidence = torch.max(probs).item() * 100

    # This model outputs: 0 = FAKE, 1 = REAL
    label = text_model.config.id2label[prediction]
    return label, confidence
# # ---------------------------
# # PREDICT IMAGE
# # ---------------------------
# def predict_image(img, image_model):
#     transform = transforms.Compose([
#         transforms.Resize((224, 224)),
#         transforms.ToTensor(),
#         transforms.Normalize([0.485, 0.456, 0.406],
#                              [0.229, 0.224, 0.225])
#     ])
#     tensor = transform(img).unsqueeze(0)
#     with torch.no_grad():
#         output = torch.sigmoid(image_model(tensor))
#     confidence = output.item() * 100
#     return confidence

# ---------------------------
# TABS
# ---------------------------
tab1, tab2 = st.tabs(["📰 Fake News Detection", "🖼️ AI Image Detection"])

with tab1:
    st.subheader("📰 Fake News Detector")

    with st.spinner("Loading text model... ⏳"):
        tokenizer, text_model = load_text_model()

    user_input = st.text_area("Paste news article here...")

    if st.button("🔍 Analyze News"):
        if user_input.strip() == "":
            st.warning("Please enter some text")
        else:
            label, conf = predict_news(user_input, tokenizer, text_model)

            if label == "FAKE":
                st.error(f"⚠️ FAKE NEWS ({conf:.2f}%)")
            else:
                st.success(f"✅ REAL NEWS ({conf:.2f}%)")

            st.progress(int(conf))

with tab2:
    st.subheader("🖼️ AI Image Detector")

    with st.spinner("Loading image model... ⏳"):
        image_model = load_image_model()

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

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

        label, confidence = predict_image(img, image_model)

        if label == "AI":
            st.error(f"⚠️ AI GENERATED IMAGE ({confidence:.2f}%)")
        else:
            st.success(f"✅ REAL IMAGE ({confidence:.2f}%)")

        st.progress(int(confidence))

# ---------------------------
# FOOTER
# ---------------------------
st.markdown("---")
st.caption("Built with ❤️ using Transformers & Deep Learning")