import gradio as gr
from rag_engine import RAGEngine
import torch
import os
import logging
import traceback

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def safe_search(query, max_results):
    """Wrapper function to handle errors gracefully"""
    try:
        rag = RAGEngine()
        results = rag.search_and_process(query, max_results)
        
        if 'error' in results:
            return f"# ❌ Error\nSorry, an error occurred while processing your search:\n```\n{results['error']}\n```"
            
        return format_results(results)
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        logger.error(error_msg)
        return f"# ❌ Error\nSorry, an error occurred while processing your search:\n```\n{str(e)}\n```"

def format_results(results):
    """Format search results into a clean markdown output"""
    if 'error' in results:
        return f"❌ Error: {results['error']}"
        
    output = []
    
    # Add insights section
    if 'insights' in results:
        insights = results['insights']
        output.append("# 🎯 Key Insights\n")
        
        if 'summary' in insights:
            output.append(insights['summary'])
            output.append("\n")
            
        if 'key_points' in insights and len(insights['key_points']) > 5:
            output.append("\n## 📌 Additional Points\n")
            for point in insights['key_points'][5:]:
                output.append(f"• {point}")
            output.append("\n")
    
    # Add sources section
    if 'insights' in results and 'sources' in results['insights']:
        output.append("\n# 📚 Sources\n")
        for idx, source in enumerate(results['insights']['sources'], 1):
            output.append(f"\n## {idx}. {source['title']}\n")
            if 'url' in source:
                output.append(f"🔗 [View Source]({source['url']})\n")
            if 'summary' in source:
                output.append(f"\n{source['summary']}\n")
    
    # Add follow-up questions
    if 'follow_up_questions' in results:
        output.append("\n# ❓ Suggested Questions\n")
        for question in results['follow_up_questions']:
            output.append(f"• {question}\n")
    
    return "\n".join(output)

def create_demo():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="Web Search + RAG") as demo:
        gr.Markdown("# 🔍 Intelligent Web Search")
        gr.Markdown("Search the web with AI-powered insights and analysis.")
        
        with gr.Row():
            with gr.Column():
                query = gr.Textbox(
                    label="Search Query",
                    placeholder="Enter your search query...",
                    lines=2
                )
                max_results = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=1,
                    label="Number of Results"
                )
                search_button = gr.Button("🔍 Search")
            
        output = gr.Markdown(
            label="Search Results",
            show_label=True
        )
        
        search_button.click(
            fn=safe_search,
            inputs=[query, max_results],
            outputs=output
        )
        
        gr.Examples(
            examples=[
                ["What is RAG in AI?", 5],
                ["Latest developments in quantum computing", 3],
                ["How does BERT work?", 5]
            ],
            inputs=[query, max_results]
        )
        
    return demo

# Create the demo
demo = create_demo()

# Launch for Spaces
demo.launch()