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):
    start_time = time.time()
    
    # Search Google
    urls = search_google(question, num_results=3)
    if not urls:
        return "⚠️ No search results found. Try a different query."
    
    # Scrape URLs
    texts_images = []
    for url in urls:
        texts_images.append(scrape_url(url))
        time.sleep(0.5)  # Add delay between requests
    
    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 only if we have texts
    if texts:
        vs.add_texts(texts)
    
    # Retrieve context
    relevant = vs.retrieve(question, top_k=2) if vs.has_data() else []
    context = "\n\n".join(relevant) if relevant else "No relevant context found."
    
    # Generate answer
    answer = generate_answer(context, question)
    
    # Prepare output
    image_markdown = ""
    for i, (url, imgs) in enumerate(zip(urls, images)):
        if imgs:
            # Show first image with source link
            img_url = imgs[0]
            image_markdown += f'<div style="margin-bottom: 20px;">'
            image_markdown += f'<a href="{url}" target="_blank"><img src="{img_url}" style="max-width: 300px; max-height: 200px;"></a><br>'
            image_markdown += f'<a href="{url}" target="_blank">Source {i+1}</a>'
            image_markdown += f'</div>'
    
    processing_time = round(time.time() - start_time, 2)
    final_output = f"""
    ## 🧠 Answer
    
    {answer}
    
    ---
    ## 📸 Images & Sources
    
    {image_markdown if image_markdown else "No images found"}
    
    <div style="margin-top: 20px; color: #666; font-size: 0.9em;">
        Processed in {processing_time} seconds | {len(urls)} sources searched
    </div>
    """
    return final_output

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="violet"),
    css="""
    .gradio-container {max-width: 800px !important}
    .message {padding: 10px; border-radius: 5px; margin: 10px 0;}
    .error {background: #ffebee; color: #b71c1c;}
    .warning {background: #fff8e1; color: #ff8f00;}
    """
) as demo:
    gr.Markdown("""
    # 🔍 **AI Web Research Agent**
    *Ask me anything - I'll search the web, analyze content, and provide answers with sources!*
    """)
    
    with gr.Row():
        inp = gr.Textbox(label="Your question", placeholder="e.g., Best laptop under 50,000 INR", scale=4)
        btn = gr.Button("Search", variant="primary", scale=1)
    
    out = gr.Markdown()
    
    btn.click(fn=ask_agent, inputs=inp, outputs=out)

demo.launch()