"""
Report generator module for the intelligent research system.

This module provides functionality to generate reports from search results
by scraping documents, storing them in a database, and synthesizing them
into a comprehensive report.
"""

import os
import asyncio
import logging
from typing import Dict, List, Any, Optional, Tuple, Union

from report.database.db_manager import get_db_manager, initialize_database
from report.document_scraper import get_document_scraper
from report.document_processor import get_document_processor
from report.report_synthesis import get_report_synthesizer
from report.report_detail_levels import get_report_detail_level_manager, DetailLevel

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

class ReportGenerator:
    """
    Report generator for the intelligent research system.
    
    This class provides methods to generate reports from search results
    by scraping documents, storing them in a database, and synthesizing them
    into a comprehensive report.
    """
    
    def __init__(self):
        """Initialize the report generator."""
        self.db_manager = get_db_manager()
        self.document_scraper = get_document_scraper()
        self.document_processor = get_document_processor()
        self.report_synthesizer = get_report_synthesizer()
        self.detail_level_manager = get_report_detail_level_manager()
        self.detail_level = "standard"  # Default detail level
        self.model_name = None  # Will use default model based on detail level
    
    async def initialize(self):
        """Initialize the report generator by setting up the database."""
        await initialize_database()
        logger.info("Report generator initialized")
    
    def set_detail_level(self, detail_level: str) -> None:
        """
        Set the detail level for report generation.
        
        Args:
            detail_level: Detail level (brief, standard, detailed, comprehensive)
        """
        try:
            # Validate detail level
            config = self.detail_level_manager.get_detail_level_config(detail_level)
            self.detail_level = detail_level
            
            # Update model if needed
            model = config.get("model")
            if model and model != self.model_name:
                self.model_name = model
                self.report_synthesizer = get_report_synthesizer(model)
            
            logger.info(f"Detail level set to {detail_level} with model {model}")
        except ValueError as e:
            logger.error(f"Error setting detail level: {e}")
            raise
    
    def get_detail_level_config(self) -> Dict[str, Any]:
        """
        Get the current detail level configuration.
        
        Returns:
            Dictionary of configuration parameters for the current detail level
        """
        return self.detail_level_manager.get_detail_level_config(self.detail_level)
    
    def get_available_detail_levels(self) -> List[Tuple[str, str]]:
        """
        Get a list of available detail levels with descriptions.
        
        Returns:
            List of tuples containing detail level and description
        """
        return self.detail_level_manager.get_available_detail_levels()

    async def process_search_results(self, search_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Process search results by scraping the URLs and storing them in the database.
        
        Args:
            search_results: List of search results, each containing at least a 'url' field
            
        Returns:
            List of processed documents
        """
        # Extract URLs from search results
        urls = [result.get('url') for result in search_results if result.get('url')]
        
        # Extract relevance scores if available
        relevance_scores = {}
        for result in search_results:
            if result.get('url') and result.get('score') is not None:
                relevance_scores[result.get('url')] = result.get('score')
        
        # Scrape URLs and store in database
        documents = await self.document_scraper.scrape_urls(urls)
        
        # Log results
        logger.info(f"Processed {len(documents)} documents out of {len(urls)} URLs")
        
        return documents, relevance_scores
    
    async def get_document_by_url(self, url: str) -> Optional[Dict[str, Any]]:
        """
        Get a document by its URL.
        
        Args:
            url: URL of the document
            
        Returns:
            Document as a dictionary, or None if not found
        """
        return await self.db_manager.get_document_by_url(url)
    
    async def search_documents(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Search for documents in the database.
        
        Args:
            query: Search query
            limit: Maximum number of results to return
            
        Returns:
            List of matching documents
        """
        return await self.db_manager.search_documents(query, limit)
    
    async def prepare_documents_for_report(self, 
                                         search_results: List[Dict[str, Any]], 
                                         token_budget: Optional[int] = None,
                                         chunk_size: Optional[int] = None,
                                         overlap_size: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Prepare documents for report generation by processing search results,
        prioritizing documents, and chunking them to fit within token budget.
        
        Args:
            search_results: List of search results
            token_budget: Maximum number of tokens to use
            chunk_size: Maximum number of tokens per chunk
            overlap_size: Number of tokens to overlap between chunks
            
        Returns:
            List of selected document chunks
        """
        # Get configuration from detail level if not specified
        config = self.get_detail_level_config()
        
        if token_budget is None:
            token_budget = config.get("token_budget")
        
        if chunk_size is None:
            chunk_size = config.get("chunk_size", 1000)
        
        if overlap_size is None:
            overlap_size = config.get("overlap_size", 100)
        
        logger.info(f"Preparing documents with token_budget={token_budget}, chunk_size={chunk_size}, overlap_size={overlap_size}")
        
        # Process search results to get documents and relevance scores
        documents, relevance_scores = await self.process_search_results(search_results)
        
        # Prioritize and chunk documents
        selected_chunks = self.document_processor.process_documents_for_report(
            documents,
            relevance_scores,
            token_budget,
            chunk_size,
            overlap_size
        )
        
        return selected_chunks
    
    async def generate_report(self, 
                             search_results: List[Dict[str, Any]],
                             query: str,
                             token_budget: Optional[int] = None,
                             chunk_size: Optional[int] = None,
                             overlap_size: Optional[int] = None,
                             detail_level: Optional[str] = None) -> str:
        """
        Generate a report from search results.
        
        Args:
            search_results: List of search results
            query: Original search query
            token_budget: Maximum number of tokens to use
            chunk_size: Maximum number of tokens per chunk
            overlap_size: Number of tokens to overlap between chunks
            detail_level: Level of detail for the report (brief, standard, detailed, comprehensive)
            
        Returns:
            Generated report as a string
        """
        # Set detail level if specified
        if detail_level:
            self.set_detail_level(detail_level)
        
        # Prepare documents for report
        selected_chunks = await self.prepare_documents_for_report(
            search_results,
            token_budget,
            chunk_size,
            overlap_size
        )
        
        # Generate report using report synthesizer
        report = await self.report_synthesizer.synthesize_report(
            selected_chunks, 
            query,
            detail_level=self.detail_level
        )
        
        return report


# Create a singleton instance for global use
report_generator = ReportGenerator()

async def initialize_report_generator():
    """Initialize the report generator."""
    await report_generator.initialize()

def get_report_generator() -> ReportGenerator:
    """
    Get the global report generator instance.
    
    Returns:
        ReportGenerator instance
    """
    return report_generator

async def test_report_generator(use_mock: bool = False):
    """
    Test the report generator with sample search results.
    
    Args:
        use_mock: If True, use mock data instead of making actual API calls
    """
    # Initialize the report generator
    await initialize_report_generator()
    
    # Get document scraper with mock option
    document_scraper = get_document_scraper(use_mock=use_mock)
    
    # Sample search results with real, accessible URLs
    search_results = [
        {
            'title': 'Python Documentation',
            'url': 'https://docs.python.org/3/tutorial/index.html',
            'snippet': 'The Python Tutorial.',
            'score': 0.95
        },
        {
            'title': 'Python Requests Library',
            'url': 'https://requests.readthedocs.io/en/latest/',
            'snippet': 'Requests is an elegant and simple HTTP library for Python.',
            'score': 0.85
        },
        {
            'title': 'Real Python',
            'url': 'https://realpython.com/',
            'snippet': 'Python tutorials for developers of all skill levels.',
            'score': 0.75
        }
    ]
    
    try:
        # Process search results
        documents, relevance_scores = await report_generator.process_search_results(search_results)
        
        # Print documents
        print(f"Processed {len(documents)} documents")
        for doc in documents:
            print(f"Document: {doc.get('title')} ({doc.get('url')})")
            print(f"Token count: {doc.get('token_count')}")
            content_preview = doc.get('content', '')[:100] + '...' if doc.get('content') else 'No content'
            print(f"Content snippet: {content_preview}")
            print()
        
        # Generate report
        report = await report_generator.generate_report(search_results, "Python programming")
        
        # Print report
        print("Generated Report:")
        print(report)
    except Exception as e:
        logger.error(f"Error during report generation test: {str(e)}")
        import traceback
        traceback.print_exc()

# Run test if this module is executed directly
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Test the report generator')
    parser.add_argument('--mock', action='store_true', help='Use mock data instead of making actual API calls')
    args = parser.parse_args()
    
    print(f"Running test with {'mock data' if args.mock else 'real data'}")
    asyncio.run(test_report_generator(use_mock=args.mock))
