# app.py
import gradio as gr
import logging
from modules.input_handler import InputHandler
from modules.retriever import Retriever
from modules.analyzer import Analyzer
from modules.citation import CitationManager
from modules.formatter import OutputFormatter
import os

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

class ResearchOrchestrator:
    def __init__(self, input_handler, retriever, analyzer, citation_manager, formatter):
        self.input_handler = input_handler
        self.retriever = retriever
        self.analyzer = analyzer
        self.citation_manager = citation_manager
        self.formatter = formatter
    
    def run(self, query):
        """Execute the research pipeline with streaming updates"""
        try:
            logging.info(f"Starting research for query: {query}")
            
            # Step 1: Process input
            yield "🔍 Processing your query..."
            processed_query = self.input_handler.process_query(query)
            logging.info("Query processed successfully")
            
            # Step 2: Retrieve data
            yield "🌐 Searching for relevant information..."
            search_results = self.retriever.search(processed_query)
            
            if not search_results:
                yield "⚠️ No relevant information found for your query. Please try rephrasing."
                logging.warning("No search results found")
                return
            
            logging.info(f"Retrieved {len(search_results)} results")
            
            # Step 3: Analyze content
            yield "🧠 Analyzing search results..."
            analysis = self.analyzer.analyze(query, search_results)
            logging.info("Analysis completed")
            
            # Step 4: Manage citations
            yield "📎 Adding citations..."
            cited_analysis = self.citation_manager.add_citations(analysis, search_results)
            logging.info("Citations added")
            
            # Step 5: Format output
            yield "✨ Formatting response..."
            formatted_output = self.formatter.format_response(cited_analysis, search_results)
            logging.info("Response formatted successfully")
            
            yield formatted_output
            
        except Exception as e:
            error_msg = f"❌ An error occurred: {str(e)}"
            logging.error(f"Error in research pipeline: {str(e)}", exc_info=True)
            yield error_msg

# Configuration
CONFIG = {
    "hf_api_base": "https://zxzbfrlg3ssrk7d9.us-east-1.aws.endpoints.huggingface.cloud/v1/",
    "hf_api_key": os.getenv("HF_TOKEN"),
    "tavily_api_key": os.getenv("TAVILY_API_KEY"),
}

# Initialize modules with error handling
def initialize_modules():
    """Initialize all modules with proper error handling"""
    try:
        if not CONFIG["tavily_api_key"]:
            raise ValueError("TAVILY_API_KEY environment variable is not set")
        
        if not CONFIG["hf_api_key"]:
            raise ValueError("HF_TOKEN environment variable is not set")
            
        input_handler = InputHandler()
        retriever = Retriever(api_key=CONFIG["tavily_api_key"])
        analyzer = Analyzer(base_url=CONFIG["hf_api_base"], api_key=CONFIG["hf_api_key"])
        citation_manager = CitationManager()
        formatter = OutputFormatter()
        
        return ResearchOrchestrator(
            input_handler, 
            retriever, 
            analyzer, 
            citation_manager, 
            formatter
        )
    except Exception as e:
        logging.error(f"Failed to initialize modules: {str(e)}")
        raise

# Initialize orchestrator
orchestrator = initialize_modules()

def research_assistant(query):
    """Main entry point for the research assistant with streaming"""
    logging.info(f"Research assistant called with query: {query}")
    for step in orchestrator.run(query):
        yield step

# Create Gradio interface
with gr.Blocks(title="Research Assistant") as demo:
    gr.Markdown("# 🧠 AI Research Assistant")
    gr.Markdown("Enter a research topic to get a structured analysis with sources")
    
    with gr.Row():
        with gr.Column():
            query_input = gr.Textbox(
                label="Research Query",
                placeholder="Enter your research question...",
                lines=3
            )
            submit_btn = gr.Button("Research", variant="primary")
        
        with gr.Column():
            output = gr.Markdown(label="Analysis Results")
    
    examples = gr.Examples(
        examples=[
            "Latest advancements in quantum computing",
            "Impact of climate change on global agriculture",
            "Recent developments in Alzheimer's treatment research"
        ],
        inputs=query_input
    )
    
    submit_btn.click(
        fn=research_assistant,
        inputs=query_input,
        outputs=output
    )
    
    query_input.submit(
        fn=research_assistant,
        inputs=query_input,
        outputs=output
    )

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