import gradio as gr
from openai import OpenAI
from smolagents import DuckDuckGoSearchTool
import re
import time

web_search = DuckDuckGoSearchTool()

SYSTEM_PROMPT = """
You are an AI research assistant that can search the web. Follow these steps:

1. FIRST ANALYZE the user's question:
   - If information is missing or ambiguous, ask ONE clarifying question
   - If clear, proceed to search

2. When searching:
   - Generate multiple specific search queries wrapped in <search> tags
   - Focus on factual keywords, one query per line
   Example:
   <search>
   Pont des Arts exact length meters
   History of Pont des Arts bridge
   </search>

3. After receiving results:
   - Analyze information from multiple sources
   - Cross-verify facts
   - If needed, generate follow-up searches
   - Provide final answer with:
     - Clear structure
     - Key facts with sources
     - Concise explanations

Never invent information. Cite sources for all facts. Use neutral, academic tone.
"""

def process_searches(response):
    searches = re.findall(r'<search>(.*?)</search>', response, re.DOTALL)
    if searches:
        return [q.strip() for q in searches[0].split('\n') if q.strip()]
    return None

def search_with_retry(query, max_retries=3, delay=2):
    for attempt in range(max_retries):
        try:
            return web_search(query)
        except Exception as e:
            if attempt < max_retries - 1:
                time.sleep(delay)
                continue
            raise
    return None

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    openrouter_key,
):
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=openrouter_key,
    )
    
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    
    full_response = ""
    search_cycle = True
    
    try:
        while search_cycle:
            search_cycle = False
            
            try:
                completion = client.chat.completions.create(
                    model="qwen/qwq-32b:free",
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    stream=True,
                    extra_headers={
                        "HTTP-Referer": "https://your-domain.com",
                        "X-Title": "Web Research Agent"
                    }
                )
            except Exception as e:
                yield f"⚠️ API Error: {str(e)}"
                return

            response = ""
            for chunk in completion:
                token = chunk.choices[0].delta.content or ""
                response += token
                full_response += token
                yield full_response  # Direct streaming without processing

            queries = process_searches(response)
            
            if queries:
                search_cycle = True
                messages.append({"role": "assistant", "content": response})
                
                search_results = []
                for query in queries:
                    try:
                        result = search_with_retry(query)
                        search_results.append(f"Search results for '{query}':\n{result}")
                    except Exception as e:
                        search_results.append(f"Search failed for '{query}': {str(e)}")
                        time.sleep(2)
                
                messages.append({
                    "role": "user",
                    "content": f"SEARCH RESULTS:\n{chr(10).join(search_results)}"
                })
                yield full_response + "\n🔍 Analyzing results...\n"

    except Exception as e:
        yield f"⚠️ Error: {str(e)}"

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", lines=8),
        gr.Slider(1000, 15000, 6000, step=500, label="Max Tokens"),
        gr.Slider(0.1, 1.0, 0.5, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, 0.85, step=0.05, label="Top-p"),
        gr.Textbox(label="OpenRouter Key", type="password")
    ],
    title="Web Research Agent",
    description="AI assistant with web search capabilities",
    examples=[
        ["Compare COVID-19 mortality rates between US and Sweden with sources"],
        ["What's the current consensus on dark matter composition?"]
    ],
    cache_examples=False
)

if __name__ == "__main__":
    demo.launch()