import gradio as gr
from search import search_google
from scraper import scrape_url
from rag import VectorStore
from llm import generate_answer
import time

vs = VectorStore()

def ask_agent(question):
    # Search Google
    with gr.Blocks(analytics_enabled=False) as progress_section:
        with gr.Row():
            gr.Textbox("Searching web...", show_label=False)
    
    urls = [u for u in search_google(question, num_results=3) if u.startswith("http")]
    if not urls:
        return "⚠️ No search results found. Try a different query."
    
    # Scrape URLs
    progress_section.children[0].children[0].value = "Scraping content..."
    texts_images = []
    for url in urls:
        texts_images.append(scrape_url(url))
    
    texts = [ti[0] for ti in texts_images if not ti[0].startswith("[Error")]
    images = [ti[1] for ti in texts_images]
    
    # Add to vector store
    if texts:
        vs.add_texts(texts)
    
    # Retrieve context
    progress_section.children[0].children[0].value = "Analyzing content..."
    relevant = vs.retrieve(question, top_k=2)
    context = "\n\n".join(relevant) if relevant else "No relevant context found."
    
    # Generate answer
    progress_section.children[0].children[0].value = "Generating answer..."
    answer = generate_answer(context, question)
    
    # Prepare image gallery
    image_gallery = []
    for url, imgs in zip(urls, images):
        if imgs:
            image_gallery.extend(imgs[:3])  # Show max 3 images per site
    
    # Prepare sources
    sources = "\n".join([f"- [{url}]({url})" for url in urls])
    
    return answer, image_gallery, sources

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="violet",
        font=[gr.themes.GoogleFont("Poppins")]
    ),
    css=".gradio-container {max-width: 900px !important}"
) as demo:
    gr.Markdown("""
    # 🌐 **Smart Web Research Agent**
    Ask anything - I'll search the web, analyze content, and provide answers with sources!
    """)
    
    with gr.Row():
        question = gr.Textbox(
            label="Your question", 
            placeholder="e.g., Best budget laptop 2024?",
            scale=4
        )
        submit_btn = gr.Button("Search", variant="primary", scale=1)
    
    progress = gr.Textbox(visible=False)
    
    with gr.Accordion("Answer", open=True):
        answer = gr.Markdown()
    
    with gr.Accordion("Sources", open=False):
        sources = gr.Markdown()
    
    with gr.Accordion("Images", open=False):
        gallery = gr.Gallery(
            columns=3, 
            object_fit="contain",
            height="auto"
        )
    
    submit_btn.click(
        fn=ask_agent,
        inputs=question,
        outputs=[answer, gallery, sources],
        api_name="search"
    )

demo.launch()