"""
Test script for the progressive report generation functionality.

This script tests the progressive report generation approach for comprehensive reports.
"""

import os
import sys
import asyncio
import logging
from typing import Dict, List, Any, Optional

# Add the project root directory to the Python path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from report.progressive_report_synthesis import get_progressive_report_synthesizer
from report.report_generator import get_report_generator, initialize_report_generator
from report.report_detail_levels import get_report_detail_level_manager

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

# Sample document chunks for testing
SAMPLE_CHUNKS = [
    {
        "document_id": "1",
        "title": "Introduction to Electric Vehicles",
        "url": "https://example.com/ev-intro",
        "content": """
        Electric vehicles (EVs) are automobiles that are propelled by one or more electric motors, using energy stored in rechargeable batteries. Compared to internal combustion engine (ICE) vehicles, EVs are quieter, have no exhaust emissions, and lower emissions overall. In the long run, EVs are often cheaper to maintain due to fewer moving parts and the increasing efficiency of battery technology.
        
        The first practical production EVs were produced in the 1880s. However, internal combustion engines were preferred for road vehicles for most of the 20th century. EVs saw a resurgence in the 21st century due to technological developments, and an increased focus on renewable energy and potential reduction of transportation's impact on climate change and other environmental issues.
        """,
        "priority_score": 0.95
    },
    {
        "document_id": "2",
        "title": "Environmental Impact of Electric Vehicles",
        "url": "https://example.com/ev-environment",
        "content": """
        The environmental impact of electric vehicles (EVs) is a complex topic that requires consideration of multiple factors. While EVs produce zero direct emissions, their overall environmental impact depends on how the electricity used to charge them is generated.
        
        In regions where electricity is produced from low-carbon sources like renewables or nuclear, EVs offer significant environmental benefits over conventional vehicles. However, in areas heavily dependent on coal or other fossil fuels for electricity generation, the benefits may be reduced.
        
        Life cycle assessments show that EVs typically have a higher environmental impact during manufacturing, primarily due to battery production, but this is usually offset by lower emissions during operation. The total lifecycle emissions of an EV are generally lower than those of a comparable conventional vehicle, especially as the vehicle is used over time.
        """,
        "priority_score": 0.9
    },
    {
        "document_id": "3",
        "title": "Economic Considerations of Electric Vehicles",
        "url": "https://example.com/ev-economics",
        "content": """
        The economics of electric vehicles (EVs) involve several factors including purchase price, operating costs, maintenance, and resale value. While EVs typically have higher upfront costs compared to conventional vehicles, they often have lower operating and maintenance costs.
        
        The total cost of ownership (TCO) analysis shows that EVs can be economically competitive or even advantageous over the vehicle's lifetime, especially in regions with high fuel prices or significant incentives for EV adoption. Factors affecting TCO include:
        
        1. Purchase price and available incentives
        2. Electricity costs versus fuel costs
        3. Maintenance requirements and costs
        4. Battery longevity and replacement costs
        5. Resale value
        
        Government incentives, including tax credits, rebates, and other benefits, can significantly reduce the effective purchase price of EVs, making them more competitive with conventional vehicles.
        """,
        "priority_score": 0.85
    },
    {
        "document_id": "4",
        "title": "Electric Vehicle Battery Technology",
        "url": "https://example.com/ev-batteries",
        "content": """
        Battery technology is a critical component of electric vehicles (EVs). Most modern EVs use lithium-ion batteries, which offer high energy density, low self-discharge, and no memory effect. However, these batteries face challenges including limited range, long charging times, degradation over time, and resource constraints for materials like lithium, cobalt, and nickel.
        
        Research and development in battery technology focus on several areas:
        
        1. Increasing energy density to improve range
        2. Reducing charging time through fast-charging technologies
        3. Extending battery lifespan and reducing degradation
        4. Developing batteries with more abundant and sustainable materials
        5. Improving safety and thermal management
        
        Solid-state batteries represent a promising future technology, potentially offering higher energy density, faster charging, longer lifespan, and improved safety compared to current lithium-ion batteries.
        """,
        "priority_score": 0.8
    },
    {
        "document_id": "5",
        "title": "Electric Vehicle Infrastructure",
        "url": "https://example.com/ev-infrastructure",
        "content": """
        Electric vehicle (EV) infrastructure refers to the charging stations, grid capacity, and supporting systems necessary for widespread EV adoption. The availability and accessibility of charging infrastructure is a critical factor in EV adoption rates.
        
        Charging infrastructure can be categorized into three main types:
        
        1. Level 1 (120V AC): Standard household outlet, providing about 2-5 miles of range per hour of charging
        2. Level 2 (240V AC): Dedicated charging station providing about 10-30 miles of range per hour
        3. DC Fast Charging: High-powered stations providing 60-80% charge in 20-30 minutes
        
        The development of EV infrastructure faces several challenges, including:
        
        - High installation costs, particularly for fast-charging stations
        - Grid capacity constraints in areas with high EV adoption
        - Standardization of charging connectors and protocols
        - Equitable distribution of charging infrastructure
        
        Government initiatives, utility programs, and private investments are all contributing to the expansion of EV charging infrastructure globally.
        """,
        "priority_score": 0.75
    },
    {
        "document_id": "6",
        "title": "Future Trends in Electric Vehicles",
        "url": "https://example.com/ev-future",
        "content": """
        The electric vehicle (EV) market is rapidly evolving, with several key trends shaping its future:
        
        1. Increasing range: Newer EV models are offering ranges exceeding 300 miles on a single charge, addressing one of the primary concerns of potential adopters.
        
        2. Decreasing battery costs: Battery costs have declined by approximately 85% since 2010, making EVs increasingly cost-competitive with conventional vehicles.
        
        3. Autonomous driving features: Many EVs are at the forefront of autonomous driving technology, with features like advanced driver assistance systems (ADAS) becoming more common.
        
        4. Vehicle-to-grid (V2G) technology: This allows EVs to not only consume electricity but also return it to the grid during peak demand, potentially creating new economic opportunities for EV owners.
        
        5. Wireless charging: Development of inductive charging technology could eliminate the need for physical connections to charge EVs.
        
        6. Integration with renewable energy: Synergies between EVs and renewable energy sources like solar and wind power are being explored to create more sustainable transportation systems.
        
        These trends suggest that EVs will continue to gain market share and could potentially become the dominant form of personal transportation in many markets within the next few decades.
        """,
        "priority_score": 0.7
    }
]

async def test_progressive_report_generation():
    """Test the progressive report generation functionality."""
    # Initialize the report generator
    await initialize_report_generator()
    
    # Get the progressive report synthesizer
    synthesizer = get_progressive_report_synthesizer()
    
    # Define a progress callback
    def progress_callback(progress, total, current_report):
        logger.info(f"Progress: {progress:.2%} ({total} chunks)")
    
    # Set progress callback
    synthesizer.set_progress_callback(progress_callback)
    
    # Test query
    query = "What are the environmental and economic impacts of electric vehicles?"
    
    logger.info(f"Starting progressive report generation for query: {query}")
    
    # Generate report progressively
    report = await synthesizer.synthesize_report_progressively(
        SAMPLE_CHUNKS, 
        query, 
        query_type="comparative", 
        detail_level="comprehensive"
    )
    
    # Print report state
    logger.info(f"Report generation completed after {synthesizer.report_state.version} iterations")
    logger.info(f"Processed {len(synthesizer.report_state.processed_chunks)} chunks")
    logger.info(f"Improvement scores: {synthesizer.report_state.improvement_scores}")
    logger.info(f"Termination reason: {synthesizer.report_state.termination_reason}")
    
    # Save the report to a file
    with open("progressive_report_test_output.md", "w") as f:
        f.write(report)
    
    logger.info(f"Report saved to progressive_report_test_output.md")
    
    return report

async def test_report_generator_with_progressive_synthesis():
    """Test the report generator with progressive synthesis for comprehensive detail level."""
    # Initialize the report generator
    await initialize_report_generator()
    
    # Get the report generator
    generator = get_report_generator()
    
    # Set detail level to comprehensive
    generator.set_detail_level("comprehensive")
    
    # Create mock search results
    search_results = [
        {
            'title': chunk['title'],
            'url': chunk['url'],
            'snippet': chunk['content'][:100] + '...',
            'score': chunk['priority_score']
        }
        for chunk in SAMPLE_CHUNKS
    ]
    
    # Test query
    query = "What are the environmental and economic impacts of electric vehicles?"
    
    logger.info(f"Starting report generation with progressive synthesis for query: {query}")
    
    # Generate report
    report = await generator.generate_report(search_results, query)
    
    # Save the report to a file
    with open("report_generator_progressive_test_output.md", "w") as f:
        f.write(report)
    
    logger.info(f"Report saved to report_generator_progressive_test_output.md")
    
    return report

async def compare_progressive_vs_standard():
    """Compare progressive synthesis with standard map-reduce approach."""
    # Initialize the report generator
    await initialize_report_generator()
    
    # Get the synthesizers
    progressive_synthesizer = get_progressive_report_synthesizer()
    standard_synthesizer = get_progressive_report_synthesizer()  # Using the same class but different method
    
    # Test query
    query = "What are the environmental and economic impacts of electric vehicles?"
    
    logger.info("Starting comparison between progressive and standard synthesis")
    
    # Generate report using progressive synthesis
    logger.info("Generating report with progressive synthesis...")
    progressive_start_time = asyncio.get_event_loop().time()
    progressive_report = await progressive_synthesizer.synthesize_report_progressively(
        SAMPLE_CHUNKS, 
        query, 
        query_type="comparative", 
        detail_level="comprehensive"
    )
    progressive_end_time = asyncio.get_event_loop().time()
    progressive_duration = progressive_end_time - progressive_start_time
    
    # Generate report using standard map-reduce
    logger.info("Generating report with standard map-reduce...")
    standard_start_time = asyncio.get_event_loop().time()
    standard_report = await standard_synthesizer.synthesize_report(
        SAMPLE_CHUNKS, 
        query, 
        query_type="comparative", 
        detail_level="detailed"  # Using detailed instead of comprehensive to use map-reduce
    )
    standard_end_time = asyncio.get_event_loop().time()
    standard_duration = standard_end_time - standard_start_time
    
    # Save reports to files
    with open("progressive_synthesis_report.md", "w") as f:
        f.write(progressive_report)
    
    with open("standard_synthesis_report.md", "w") as f:
        f.write(standard_report)
    
    # Compare results
    logger.info(f"Progressive synthesis took {progressive_duration:.2f} seconds")
    logger.info(f"Standard synthesis took {standard_duration:.2f} seconds")
    logger.info(f"Progressive report length: {len(progressive_report)} characters")
    logger.info(f"Standard report length: {len(standard_report)} characters")
    
    return {
        "progressive": {
            "duration": progressive_duration,
            "length": len(progressive_report),
            "iterations": progressive_synthesizer.report_state.version
        },
        "standard": {
            "duration": standard_duration,
            "length": len(standard_report)
        }
    }

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Test progressive report generation')
    parser.add_argument('--test', choices=['progressive', 'generator', 'compare'], default='progressive',
                        help='Test to run (progressive, generator, or compare)')
    args = parser.parse_args()
    
    if args.test == 'progressive':
        asyncio.run(test_progressive_report_generation())
    elif args.test == 'generator':
        asyncio.run(test_report_generator_with_progressive_synthesis())
    elif args.test == 'compare':
        asyncio.run(compare_progressive_vs_standard())
