"""qResearch: Dual-Agent Research System"""

import os
import gradio as gr
from smolagents import CodeAgent, HfApiModel, tool
from typing import Dict, List, Optional, Tuple, Union
import requests
from urllib.parse import quote_plus

# Import browser components
from scripts.text_web_browser import SimpleTextBrowser, SearchInformationTool
from scripts.cookies import COOKIES

class DuckDuckGoSearchTool(SearchInformationTool):
    """Search tool that uses DuckDuckGo's HTML interface"""
    
    def forward(self, query: str, filter_year: Optional[int] = None) -> str:
        """Performs search using DuckDuckGo's HTML endpoint"""
        encoded_query = quote_plus(query)
        # Use the HTML endpoint directly
        url = f"https://html.duckduckgo.com/html/?q={encoded_query}"
        print(f"DEBUG: Searching with URL: {url}")
        self.browser.visit_page(url)
        header, content = self.browser._state()
        return header.strip() + "\n=======================\n" + content

@tool
def analyze_content(text: str, analysis_type: str = "general") -> str:
    """Analyzes content for various aspects like key points, themes, or citations
    
    Args:
        text: The content text to be analyzed for key points and themes
        analysis_type: Type of analysis to perform ("general", "academic", "citations")
    
    Returns:
        str: Structured analysis results including key points and findings
    """
    if "academic" in analysis_type.lower():
        return (
            "Academic Analysis:\n"
            "1. Main Arguments:\n"
            f"   - Key points from text: {text[:200]}...\n"
            "2. Evidence Quality:\n"
            "   - Source credibility assessment\n"
            "   - Data verification\n"
            "3. Research Context:\n"
            "   - Field relevance\n"
            "   - Current research status"
        )
    else:
        return (
            "General Analysis:\n"
            "1. Key Findings:\n"
            f"   - Main points from content: {text[:200]}...\n"
            "2. Supporting Evidence:\n"
            "   - Data and examples\n"
            "3. Practical Applications:\n"
            "   - Real-world relevance\n"
            "   - Implementation possibilities"
        )

class ResearchSystem:
    def __init__(self):
        # Initialize browser with request settings
        self.browser = SimpleTextBrowser(
            viewport_size=4096,
            downloads_folder="./downloads",
            request_kwargs={
                "cookies": COOKIES,
                "headers": {
                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
                    "Accept-Language": "en-US,en;q=0.5",
                    "DNT": "1",
                    "Connection": "keep-alive",
                    "Upgrade-Insecure-Requests": "1"
                }
            }
        )
        
        # Initialize model
        self.model = HfApiModel(
            model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
            custom_role_conversions={
                "tool-call": "assistant",
                "tool-response": "user"
            }
        )

        # Initialize agent with custom DuckDuckGo search tool
        self.researcher = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(self.browser),
                analyze_content
            ],
            model=self.model
        )
        
        self.formatter = CodeAgent(
            tools=[],
            model=self.model
        )

    def create_interface(self):
        with gr.Blocks(title="qResearch", theme=gr.themes.Soft()) as interface:
            gr.Markdown(
                "# qResearch\n"
                "*Research → Analysis*\n"
                "---"
            )
            
            with gr.Row():
                with gr.Column(scale=3):
                    chat = gr.Chatbot(
                        label="Research Process",
                        height=600,
                        show_label=True,
                        type="messages"
                    )
                    
                with gr.Column(scale=1):
                    input_box = gr.Textbox(
                        label="Research Query",
                        placeholder="Enter your research topic...",
                        lines=3
                    )
                    submit_btn = gr.Button("Search", variant="primary")
            
            submit_btn.click(
                self.process_query,
                inputs=[input_box],
                outputs=[chat]
            )
        
        return interface

    def process_query(self, query: str) -> List[Dict[str, str]]:
        """Process a research query using web search and analysis"""
        try:
            print(f"\nDEBUG: Processing query: {query}")
            
            # Get search results using the browser
            search_results = self.researcher.run(f"Search for information about: {query}")
            print(f"\nDEBUG: Search completed. Results:\n{search_results}")
            
            # Analyze the results
            analysis = self.researcher.run(f"analyze_content: {search_results}")
            
            # Format in MLA style
            format_prompt = (
                "Format this research in MLA style:\n"
                f"{search_results}\n\n"
                f"Analysis:\n{analysis}"
            )
            formatted = self.formatter.run(format_prompt)
            
            return [
                {"role": "user", "content": query},
                {"role": "assistant", "content": f"📚 Research Findings:\n{search_results}\n\n📊 Analysis:\n{analysis}"},
                {"role": "assistant", "content": f"📝 MLA Formatted:\n{formatted}"}
            ]
            
        except Exception as e:
            error_msg = f"Error during research: {str(e)}"
            print(f"DEBUG: Error occurred: {error_msg}")
            return [{"role": "assistant", "content": error_msg}]

if __name__ == "__main__":
    # Create downloads directory if it doesn't exist
    os.makedirs("./downloads", exist_ok=True)
    
    system = ResearchSystem()
    system.create_interface().launch(
        server_port=7860,
        share=True,
        show_api=False
    )