import streamlit as st
import requests
import json
import time
import base64
from PIL import Image
from io import BytesIO
from transformers import pipeline
from concurrent.futures import ThreadPoolExecutor

# --- 1. CONFIGURATION & SECRETS ---
# Add your key to Streamlit/HuggingFace Secrets as 'SERPER_API_KEY' for safety
try:
    SERPER_API_KEY = st.secrets["SERPER_API_KEY"]
except:
    # Fallback for local testing
    SERPER_API_KEY = "827d931e4257327f4bcc9a35b2995001a68635d0"

SERPER_URL = "https://google.serper.dev/images"
BG_IMAGE = "vecteezy_abstract-background-design-background-texture-design-with_18752866-1.jpg"
LOADING_VIDEO = "vecteezy_loading-bar-animation_26651030.mp4"

# --- Page Config & Browser Icon ---
st.set_page_config(
    page_title="Bouncer",
    page_icon=BG_IMAGE,  # Sets your custom background as the browser tab icon
    layout="wide"
)

# --- 2. CUSTOM BACKGROUND STYLING ---
def add_custom_style(image_file):
    with open(image_file, "rb") as image:
        encoded_string = base64.b64encode(image.read()).decode()
    st.markdown(
        f"""
        <style>
        .stApp {{
            background-image: url("data:image/jpg;base64,{encoded_string}");
            background-attachment: fixed;
            background-size: cover;
        }}
        /* Sidebar Glassmorphism effect */
        [data-testid="stSidebar"] {{
            background-color: rgba(255, 255, 255, 0.05) !important;
            backdrop-filter: blur(15px);
            border-right: 1px solid rgba(255, 255, 255, 0.1);
        }}
        /* Title shadows for readability */
        h1 {{ color: white !important; text-shadow: 2px 2px 8px #000000; }}
        </style>
        """,
        unsafe_allow_html=True
    )

try:
    add_custom_style(BG_IMAGE)
except:
    st.warning("Background image not found. Ensure the filename is correct.")

# --- 3. AI MODELS ---
@st.cache_resource
def load_models():
    style_pipe = pipeline("zero-shot-image-classification", model="google/siglip-base-patch16-224", device=-1)
    anime_ai_pipe = pipeline("image-classification", model="legekka/AI-Anime-Image-Detector-ViT", device=-1)
    photo_ai_pipe = pipeline("image-classification", model="Ateeqq/ai-vs-human-image-detector", device=-1)
    return style_pipe, anime_ai_pipe, photo_ai_pipe

style_classifier, anime_ai_detector, photo_ai_detector = load_models()

# --- 4. PROCESSING LOGIC ---
def get_score(preds, target_labels):
    for p in preds:
        if p['label'].lower() in target_labels: return p['score']
    return 0.0

def download_and_process(item, tolerance):
    url = item.get("imageUrl") 
    try:
        response = requests.get(url, timeout=5)
        # Robust opening to handle transparency
        raw_img = Image.open(BytesIO(response.content))
        img = raw_img.convert("RGBA").convert("RGB")
        img_small = img.resize((224, 224))
        
        style_results = style_classifier(img_small, candidate_labels=["anime illustration", "real photo"])
        top_style = style_results[0]['label']
        
        if top_style == "anime illustration":
            preds = anime_ai_detector(img_small)
            # Model uses 'natural' for human art
            human_score = get_score(preds, ['natural', 'human', 'real'])
        else:
            preds = photo_ai_detector(img_small)
            human_score = get_score(preds, ['human', 'real', 'natural'])
            
        if human_score >= tolerance:
            return {"img": img, "score": human_score, "url": url, "style": top_style}
    except: return None

# --- 5. SIDEBAR CONTROLS ---
with st.sidebar:
    st.title("Toggle go vurrrr")
    query = st.text_input("What are you looking for?", "Concept Art")

    # Target Count with Help
    c1, c2 = st.columns([4, 1])
    target_count = c1.slider("Results", 1, 40, 6)
    if c2.button("❓", key="h_count"):
        st.info("The number of human-made images you want to see.")

    # Strictness with Help
    s1, s2 = st.columns([4, 1])
    tolerance = s1.slider("Strictness", 0.0, 1.0, 0.5)
    if s2.button("❓", key="h_strict"):
        st.info("How sure the AI must be. 0.8+ is strict; 0.2 is loose.")

    # Threads with Help
    t1, t2 = st.columns([4, 1])
    workers = t1.slider("Threads", 2, 16, 8)
    if t2.button("❓", key="h_thread"):
        st.info("Processing speed. Higher is faster but heavier on memory.")

    start_search = st.button("Scanning the Internet for Target", use_container_width=True)

# --- 6. MAIN DISPLAY & LOADING BAR ---
st.markdown("Bouncer")
st.write("Creativity shall not yeild to none")

if start_search and query:
    start_time = time.time()
    
    # Custom Video Loading Bar
    loading_placeholder = st.empty()
    with loading_placeholder.container():
        st.video(LOADING_VIDEO, autoplay=True, loop=True, muted=True)
        st.write("### Cleaning Slop")

    # Fetch and process
    payload = json.dumps({"q": query, "num": target_count * 3})
    headers = {'X-API-KEY': SERPER_API_KEY, 'Content-Type': 'application/json'}
    response = requests.post(SERPER_URL, headers=headers, data=payload)
    raw_results = response.json().get("images", [])

    if raw_results:
        with ThreadPoolExecutor(max_workers=workers) as executor:
            results = list(executor.map(lambda r: download_and_process(r, tolerance), raw_results))
            final_images = [r for r in results if r is not None][:target_count]

        # Remove the loading video once finished
        loading_placeholder.empty()

        if final_images:
            st.success(f"Verified {len(final_images)} images in {time.time() - start_time:.1f}s")
            cols = st.columns(3)
            for i, item in enumerate(final_images):
                with cols[i % 3]:
                    st.image(item["img"], use_container_width=True)
                    st.caption(f"Human Confidence: {item['score']:.0%}")
        else:
            st.warning("No images passed the AI filter. Try lowering 'Strictness'.")