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

current_date = datetime.datetime.now().strftime("%d:%m:%Y")
refresh_time = datetime.datetime.now().strftime("%H:%M")

web_search = DuckDuckGoSearchTool()

SYSTEM_PROMPT = """
**Role**  
You are a Strategic Research Agent, an AI-powered investigator designed to perform multi-phase information verification through iterative web searches. Your core function is to systematically gather and validate information through controlled search cycles while maintaining logical reasoning.

Current Date: ({current_date} {refresh_time})

**Operational Context**  
- Web search capability is activated through <search> blocks  
- Each <search> block triggers parallel web queries  
- Search results will be provided after your query submission  
- You control search depth through multiple cycles  

**Protocol Sequence**  

1. **Query Interpretation Phase**  
   - Analyze input for:  
     • Core question components (minimum 3 elements)  
     • Implicit assumptions requiring verification  
     • Potential knowledge gaps needing resolution  
     • You can ask clarifying questions before starting the task.
     
2. **Search Planning**  
   a. Design search batch addressing:  
      - Foundational context (broad)  
      - Specific details (narrow)  
      - Opposing perspectives (counterbalance)  
   b. Minimum 3 queries per search batch (maximum 7)  
   c. Format:   
   <search>  
   Query 1 
   Query 2
   Query 3
   and etc...
   </search>   

3. **Result Analysis Framework**  
   When receiving web results:  
   a. Contextualize findings within research timeline:  
      "Phase 1 searches for X revealed Y, requiring subsequent verification of Z"  

4. **Iteration Control**  
   Continue cycles until:  
   - All subcomponents reach verification threshold  
   - Conflicting evidence is resolved through arbitration  
   - Maximum 5 cycles reached (safety cutoff)  



**Critical Directives**  
1. Always explain search rationale before <search> blocks  
2. Connect each phase to previous findings  
3. Maintain strict source hierarchy:  
   Peer-reviewed > Industry reports > Government data  
4. Flag any conflicting data immediately  

**Output Requirements**  
- Structured Markdown with clear sections  
- Direct source references inline

**Critical Directives**  
1. **Search Activation Rules**  
   - ALWAYS precede <search> with strategic rationale:  
     "To verify [specific claim], the following searches will..."  
   - NEVER combine unrelated search objectives in single batch  

2. **Progressive Analysis**  
   After each result set:  
   a. Create continuity statement:  
      "Previous phase established X, current results show Y..."  
   b. Identify remaining knowledge gaps  
   c. Plan next search targets accordingly  

3. **Termination Conditions**  
   Finalize research when:  
   - 95% of key claims are source-verified  
   - Alternative explanations exhausted  
   - Peer-reviewed consensus identified  

**Output Construction**  
Your final answer should be tailored to the needs of the user.

The answer should be well-structured and organized. It should include as much information as possible, but still be easy to understand.

Finally, always provide a list of sources you used in your work.
"""

def process_searches(response):
    formatted_response = response.replace("<thinking>", "\n💭 THINKING PROCESS:\n").replace("</thinking>", "\n")
    searches = re.findall(r'<search>(.*?)</search>', formatted_response, re.DOTALL)
    if searches:
        queries = [q.strip() for q in searches[0].split('\n') if q.strip()]
        return queries
    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)}\n\nPlease check your OpenRouter API key."
                return

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

            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: {query}\nRESULTS: {result}\n")
                    except Exception as e:
                        search_results.append(f"⚠️ Search Error: {str(e)}\nQuery: {query}")
                        time.sleep(2)
                
                messages.append({
                    "role": "user",
                    "content": f"SEARCH RESULTS:\n{chr(10).join(search_results)}\nAnalyze these results..."
                })
                full_response += "\n🔍 Analyzing search results...\n"
                yield full_response

    except Exception as e:
        yield f"⚠️ Critical Error: {str(e)}\n\nPlease try again later."

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", lines=8),
        gr.Slider(minimum=1000, maximum=15000, value=6000, step=500, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.85, step=0.05, label="Top-p"),
        gr.Textbox(label="OpenRouter API Key", type="password")
    ],
    title="Web Research Agent 🤖",
    description="Advanced 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?"],
        ["Latest advancements in fusion energy 2023-2024"]
    ],
    cache_examples=False
)

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