import gradio as gr
from PIL import Image as PILImage
import tempfile

from modules.ddg_search import search_images
from modules.embeddings import get_embedding, load_image_from_url
from modules.similarity import compute_similarity
from modules.pdf_report import create_pdf_report


def analyze_image_stream(prompt, uploaded_img, num_results):
    if not uploaded_img or not prompt:
        yield None, None, None, "❌ Please provide both a prompt and an image.", *[gr.update(visible=False)] * 6
        return

    user_img = uploaded_img.convert("RGB")
    urls = search_images(prompt, max_results=num_results)
    retrieved_images = [(url, "") for url in urls]

    yield retrieved_images, None, None, "🔄 Computing similarity...", \
          gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), \
          gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)

    user_emb = get_embedding(user_img)
    results = []
    for url in urls:
        try:
            img = load_image_from_url(url)
            emb = get_embedding(img)
            sim = compute_similarity(user_emb, emb)
            results.append((url, sim))
        except:
            continue

    results = sorted(results, key=lambda x: x[1], reverse=True)
    top_results = results[:5]
    top_images = [(url, f"Similarity: {sim:.3f}") for url, sim in top_results]

    pdf_bytes = create_pdf_report(prompt, user_img, top_results)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as f:
        f.write(pdf_bytes)
        pdf_path = f.name

    yield retrieved_images, top_images, pdf_path, "✅ Done!", \
          gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), \
          gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)


with gr.Blocks(title="Influence Analyzer Demo", css="""
#uploaded_image img {
    object-fit: contain;
    height: 250px;
    width: auto;
    border: 1px solid #ccc;
    margin-top: 5px;
}
""") as demo:
    gr.Markdown("""
### 🔍 Search-Based Influence Analyzer Demo
This tool supports the **interpretability of generative black-box models** by identifying potentially influential training data.

**How it works:**
- 🔎 Retrieves public images using DuckDuckGo based on your prompt.
- 🧠 Computes embeddings and compares them to your uploaded image to find the most similar results.

The goal is to provide a proxy for understanding which real-world images may have contributed to a generation — useful for **interpretability**, **data attribution**, and **copyright assessment**.
""")

    gr.Markdown("### 📝 Provide Your Prompt and Image")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Describe your image prompt...", lines=1)
            num_results = gr.Slider(5, 50, value=30, step=1, label="Images to retrieve")

            gr.Markdown("**✏️ Try an example**")
            examples = [
                ("a horse in the moon", "examples/moon_horse.png"),
                ("a greek philosopher drinking beer", "examples/greek_beer.png"),
            ]

            example_btns = []
            for i, (text, path) in enumerate(examples):
                btn = gr.Button(text)
                example_btns.append((btn, text, path))

        with gr.Column():
            uploaded_img = gr.Image(type="pil", label="Upload your generated image", elem_id="uploaded_image")

    run_btn = gr.Button("🚀 Run Search & Compare", variant="primary")

    for btn, prompt_text, image_path in example_btns:
        def make_loader(p=prompt_text, img=image_path):
            def _load():
                return p, PILImage.open(img)
            return _load
        btn.click(fn=make_loader(), inputs=[], outputs=[prompt, uploaded_img])

    gr.Markdown("---")
    retrieved_heading = gr.Markdown("### 🔍 Retrieved Images", visible=False)
    retrieved_gallery = gr.Gallery(label="", columns=5, height="auto", visible=False)

    similar_heading = gr.Markdown("### ✅ Top 5 Most Similar Images", visible=False)
    similar_gallery = gr.Gallery(label="", columns=5, height="auto", visible=False)

    output_heading = gr.Markdown("### 📄 Output Report", visible=False)
    pdf_out = gr.File(label="Download Similarity Report", visible=False)

    status = gr.Textbox(label="Status", interactive=False)

    run_btn.click(
        fn=analyze_image_stream,
        inputs=[prompt, uploaded_img, num_results],
        outputs=[
            retrieved_gallery,
            similar_gallery,
            pdf_out,
            status,
            retrieved_heading,
            similar_heading,
            output_heading,
            retrieved_gallery,
            similar_gallery,
            pdf_out,
        ],
    )

    gr.Markdown("---")
    gr.Markdown(
        "🧠 *This demo is part of ongoing research in interpretability for generative models. "
        "If you're interested in collaboration, feel free to get in touch.*"
    )

demo.launch()