"""
Simulation command implementation.

Implements the 'precipgen simulate' command for generating synthetic
precipitation using WGEN algorithm with validated parameters.
"""

import argparse
import json
from pathlib import Path
from typing import Optional
from datetime import datetime, timedelta

import pandas as pd
import numpy as np

from precipgen.utils import get_logger, DataLoadError
from precipgen.engines.simulation import SimulationEngine
from precipgen.config.precipgen_config import PrecipGenConfig
from precipgen.utils.directory_manager import DirectoryManager


def add_parser(subparsers) -> argparse.ArgumentParser:
    """
    Add simulate command parser to subparsers.
    
    Args:
        subparsers: ArgumentParser subparsers object
        
    Returns:
        The simulate command parser
    """
    parser = subparsers.add_parser(
        "simulate",
        help="Generate synthetic precipitation using WGEN algorithm",
        description="Generate synthetic precipitation time series using validated parameters",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  precipgen simulate --params params.json --days 3650 --seed 12345
  precipgen simulate --params params.json --days 3650 --trends --start-year 2025
  precipgen simulate --station USC00426128 --auto-detect --days 10950
  precipgen simulate --params params.json --interactive
        """
    )
    
    # Input parameter options
    parser.add_argument(
        "--params", "-p",
        type=str,
        help="Path to parameter manifest JSON file"
    )
    
    parser.add_argument(
        "--station", "-s",
        type=str,
        help="Station ID for auto-detecting parameters"
    )
    
    parser.add_argument(
        "--auto-detect",
        action="store_true",
        help="Auto-detect parameter file in station directory"
    )
    
    # Simulation options
    parser.add_argument(
        "--days", "-d",
        type=int,
        default=3650,
        help="Number of days to simulate (default: 3650, ~10 years)"
    )
    
    parser.add_argument(
        "--seed",
        type=int,
        help="Random seed for reproducible results"
    )
    
    parser.add_argument(
        "--trends", "-t",
        action="store_true",
        help="Enable trend mode for non-stationary projections"
    )
    
    parser.add_argument(
        "--start-year",
        type=int,
        default=2025,
        help="Start year for trend projections (default: 2025)"
    )
    
    # Interactive mode
    parser.add_argument(
        "--interactive", "-i",
        action="store_true",
        help="Interactive mode for step-by-step simulation"
    )
    
    parser.add_argument(
        "--status",
        action="store_true",
        help="Display current simulation status"
    )
    
    # Output options
    parser.add_argument(
        "--output-dir", "-o",
        type=str,
        help="Custom output directory (default: ./precipgen_output)"
    )
    
    parser.add_argument(
        "--output-file",
        type=str,
        help="Custom output filename (default: synthetic.csv)"
    )
    
    parser.add_argument(
        "--format",
        choices=["csv", "json", "both"],
        default="csv",
        help="Output format (default: csv)"
    )
    
    return parser


def load_parameter_manifest(params_path: str) -> dict:
    """
    Load parameter manifest from JSON file.
    
    Args:
        params_path: Path to parameter manifest JSON file
        
    Returns:
        Parameter manifest dictionary
        
    Raises:
        DataLoadError: If file cannot be loaded or parsed
    """
    try:
        params_file = Path(params_path)
        
        if not params_file.exists():
            raise DataLoadError(f"Parameter file not found: {params_path}")
        
        with open(params_file, 'r') as f:
            params = json.load(f)
        
        # Validate required fields
        required_fields = ['monthly_parameters']
        for field in required_fields:
            if field not in params:
                raise DataLoadError(f"Missing required field in parameters: {field}")
        
        # Validate monthly parameters
        monthly_params = params['monthly_parameters']
        for month in range(1, 13):
            month_str = str(month)
            if month_str not in monthly_params:
                raise DataLoadError(f"Missing parameters for month {month}")
            
            month_data = monthly_params[month_str]
            required_params = ['p_ww', 'p_wd', 'alpha', 'beta']
            for param in required_params:
                if param not in month_data:
                    raise DataLoadError(f"Missing parameter {param} for month {month}")
        
        return params
        
    except json.JSONDecodeError as e:
        raise DataLoadError(f"Invalid JSON in parameter file: {e}")
    except Exception as e:
        raise DataLoadError(f"Failed to load parameter manifest: {e}")


def auto_detect_params_file(station_id: str, base_dir: Optional[str] = None) -> Optional[str]:
    """
    Auto-detect parameter file in station directory.
    
    Args:
        station_id: Station identifier
        base_dir: Base output directory
        
    Returns:
        Path to parameter file or None if not found
    """
    dir_manager = DirectoryManager(base_dir or "./precipgen_output")
    
    # Check if station directory exists
    station_info = dir_manager.get_station_info(station_id, base_dir)
    if not station_info:
        return None
    
    # Look for parameter files in analysis directories
    analysis_dir = station_info.directory_path / "analysis"
    
    if not analysis_dir.exists():
        return None
    
    # Look for most recent parameter file
    param_files = []
    
    # Check timestamped subdirectories
    for subdir in analysis_dir.iterdir():
        if subdir.is_dir():
            params_file = subdir / "params.json"
            if params_file.exists():
                param_files.append(params_file)
    
    # Check root analysis directory
    params_file = analysis_dir / "params.json"
    if params_file.exists():
        param_files.append(params_file)
    
    if not param_files:
        return None
    
    # Return most recent file
    param_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
    return str(param_files[0])


def generate_synthetic_precipitation(params: dict, n_days: int, seed: Optional[int] = None,
                                   trends: bool = False, start_year: int = 2025) -> pd.Series:
    """
    Generate synthetic precipitation time series.
    
    Args:
        params: Parameter manifest dictionary
        n_days: Number of days to simulate
        seed: Random seed for reproducibility
        trends: Whether to apply trend projections
        start_year: Start year for simulation
        
    Returns:
        Pandas Series with synthetic precipitation data
    """
    if seed is not None:
        np.random.seed(seed)
    
    # Create date index
    start_date = datetime(start_year, 1, 1)
    dates = [start_date + timedelta(days=i) for i in range(n_days)]
    date_index = pd.DatetimeIndex(dates)
    
    # Initialize precipitation array
    precip = np.zeros(n_days)
    wet_state = False  # Start with dry day
    
    # Simulate day by day
    for i, date in enumerate(dates):
        month = date.month
        month_params = params['monthly_parameters'][str(month)]
        
        # Get parameters for this month
        p_ww = month_params['p_ww']  # Wet-to-wet transition probability
        p_wd = month_params['p_wd']  # Wet-to-dry transition probability (actually dry-to-wet)
        alpha = month_params['alpha']  # Gamma shape parameter
        beta = month_params['beta']    # Gamma scale parameter
        
        # Apply trends if enabled (simplified)
        if trends:
            year_progress = (date.year - start_year) / 30.0  # 30-year trend
            trend_factor = 1.0 + 0.1 * year_progress  # 10% increase over 30 years
            alpha *= trend_factor
        
        # Determine wet/dry state transition
        if wet_state:
            # Currently wet - probability of staying wet
            wet_state = np.random.random() < p_ww
        else:
            # Currently dry - probability of becoming wet
            wet_state = np.random.random() < p_wd
        
        # Generate precipitation amount if wet day
        if wet_state:
            # Generate from gamma distribution
            precip[i] = np.random.gamma(alpha, beta)
        else:
            precip[i] = 0.0
    
    return pd.Series(precip, index=date_index, name='precipitation')


def calculate_simulation_statistics(precip_data: pd.Series) -> dict:
    """
    Calculate comprehensive statistics for simulation output.
    
    Args:
        precip_data: Synthetic precipitation time series
        
    Returns:
        Dictionary with monthly, annual, and seasonal statistics
    """
    stats = {}
    
    # Monthly statistics
    monthly_stats = {}
    for month in range(1, 13):
        month_data = precip_data[precip_data.index.month == month]
        wet_days = month_data > 0.1
        
        monthly_stats[month] = {
            'mean_precip': float(month_data.mean()),
            'std_precip': float(month_data.std()),
            'total_precip': float(month_data.sum()),
            'wet_day_frequency': float(wet_days.mean()),
            'wet_days': int(wet_days.sum()),
            'max_daily': float(month_data.max())
        }
    
    stats['monthly_statistics'] = monthly_stats
    
    # Annual statistics
    annual_stats = {}
    for year in precip_data.index.year.unique():
        year_data = precip_data[precip_data.index.year == year]
        wet_days = year_data > 0.1
        
        # Calculate dry spell lengths
        dry_spells = []
        current_dry_spell = 0
        for is_wet in wet_days:
            if not is_wet:
                current_dry_spell += 1
            else:
                if current_dry_spell > 0:
                    dry_spells.append(current_dry_spell)
                current_dry_spell = 0
        
        annual_stats[year] = {
            'total_precip': float(year_data.sum()),
            'wet_days': int(wet_days.sum()),
            'max_daily': float(year_data.max()),
            'mean_dry_spell': float(np.mean(dry_spells)) if dry_spells else 0.0,
            'max_dry_spell': int(max(dry_spells)) if dry_spells else 0
        }
    
    stats['annual_statistics'] = annual_stats
    
    # Seasonal statistics
    seasons = {
        'winter': [12, 1, 2],
        'spring': [3, 4, 5], 
        'summer': [6, 7, 8],
        'fall': [9, 10, 11]
    }
    
    seasonal_stats = {}
    for season_name, months in seasons.items():
        season_data = precip_data[precip_data.index.month.isin(months)]
        wet_days = season_data > 0.1
        
        seasonal_stats[season_name] = {
            'mean_precip': float(season_data.mean()),
            'total_precip': float(season_data.sum()),
            'wet_day_frequency': float(wet_days.mean()),
            'max_daily': float(season_data.max())
        }
    
    stats['seasonal_statistics'] = seasonal_stats
    
    return stats


def save_simulation_output(precip_data: pd.Series, stats: dict, metadata: dict,
                          station_id: str, output_format: str, output_file: Optional[str],
                          dir_manager: DirectoryManager) -> list[str]:
    """
    Save simulation output to organized directory structure.
    
    Args:
        precip_data: Synthetic precipitation time series
        stats: Simulation statistics
        metadata: Simulation metadata
        station_id: Station identifier
        output_format: Output format (csv, json, or both)
        output_file: Custom output filename
        dir_manager: Directory manager instance
        
    Returns:
        List of saved file paths
    """
    # Create timestamped simulation directory
    sim_dir = dir_manager.create_timestamped_subdirectory(station_id, "simulation")
    
    saved_files = []
    
    # Save synthetic precipitation data
    if output_format in ["csv", "both"]:
        csv_filename = output_file or "synthetic.csv"
        csv_path = sim_dir / csv_filename
        
        # Save as CSV with date and precipitation columns
        df = pd.DataFrame({
            'date': precip_data.index.strftime('%Y-%m-%d'),
            'precipitation': precip_data.values
        })
        df.to_csv(csv_path, index=False)
        saved_files.append(str(csv_path))
    
    if output_format in ["json", "both"]:
        json_path = sim_dir / "synthetic.json"
        
        # Convert to JSON format
        json_data = {
            'metadata': metadata,
            'data': [
                {
                    'date': date.strftime('%Y-%m-%d'),
                    'precipitation': float(precip)
                }
                for date, precip in zip(precip_data.index, precip_data.values)
            ]
        }
        
        with open(json_path, 'w') as f:
            json.dump(json_data, f, indent=2)
        saved_files.append(str(json_path))
    
    # Save statistics
    stats_path = sim_dir / "monthly_stats.json"
    with open(stats_path, 'w') as f:
        json.dump(stats['monthly_statistics'], f, indent=2)
    saved_files.append(str(stats_path))
    
    annual_stats_path = sim_dir / "annual_stats.json"
    with open(annual_stats_path, 'w') as f:
        json.dump(stats['annual_statistics'], f, indent=2, default=str)
    saved_files.append(str(annual_stats_path))
    
    seasonal_stats_path = sim_dir / "seasonal_stats.json"
    with open(seasonal_stats_path, 'w') as f:
        json.dump(stats['seasonal_statistics'], f, indent=2)
    saved_files.append(str(seasonal_stats_path))
    
    # Save metadata
    metadata_path = sim_dir / "simulation_metadata.json"
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2, default=str)
    saved_files.append(str(metadata_path))
    
    return saved_files


def interactive_simulation(params: dict, seed: Optional[int] = None) -> None:
    """
    Run interactive step-by-step simulation.
    
    Args:
        params: Parameter manifest
        seed: Random seed
    """
    if seed is not None:
        np.random.seed(seed)
    
    print("Interactive simulation mode")
    print("Press Enter to generate next day, 'q' to quit, 's' to show status")
    
    current_date = datetime(2025, 1, 1)
    wet_state = False
    day_count = 0
    total_precip = 0.0
    
    while True:
        try:
            user_input = input(f"\nDay {day_count + 1} ({current_date.strftime('%Y-%m-%d')}): ").strip().lower()
            
            if user_input == 'q':
                break
            elif user_input == 's':
                print(f"Status: Day {day_count}, Total precip: {total_precip:.1f}mm, Current state: {'Wet' if wet_state else 'Dry'}")
                continue
            
            # Generate next day
            month = current_date.month
            month_params = params['monthly_parameters'][str(month)]
            
            p_ww = month_params['p_ww']
            p_wd = month_params['p_wd']
            alpha = month_params['alpha']
            beta = month_params['beta']
            
            # State transition
            if wet_state:
                wet_state = np.random.random() < p_ww
            else:
                wet_state = np.random.random() < p_wd
            
            # Generate precipitation
            if wet_state:
                precip = np.random.gamma(alpha, beta)
                total_precip += precip
                print(f"Wet day: {precip:.1f}mm")
            else:
                print("Dry day: 0.0mm")
            
            current_date += timedelta(days=1)
            day_count += 1
            
        except KeyboardInterrupt:
            print("\nSimulation interrupted")
            break
    
    print(f"\nSimulation summary: {day_count} days, {total_precip:.1f}mm total precipitation")


def execute(args) -> int:
    """
    Execute the simulate command.
    
    Args:
        args: Parsed command line arguments
        
    Returns:
        Exit code (0 for success, non-zero for error)
    """
    logger = get_logger("precipgen.cli.simulate")
    
    try:
        # Initialize directory manager
        dir_manager = DirectoryManager(args.output_dir or "./precipgen_output")
        
        # Determine parameter file path
        params_path = None
        
        if args.auto_detect and args.station:
            params_path = auto_detect_params_file(args.station, args.output_dir)
            if not params_path:
                print(f"Error: No parameter file found for station {args.station}")
                print("\nTroubleshooting:")
                print("  • Run analysis first to generate parameters:")
                print(f"    precipgen analyze --station {args.station} --auto-detect")
                print("  • Check that analysis completed successfully")
                print("  • Verify the station directory exists:")
                print("    precipgen list-outputs")
                return 1
            print(f"Auto-detected parameter file: {params_path}")
            
        elif args.params:
            params_path = args.params
            
        else:
            print("Error: Must specify --params file or use --station with --auto-detect")
            print("\nUsage examples:")
            print("  precipgen simulate --params params.json --days 3650")
            print("  precipgen simulate --station USC00426128 --auto-detect --days 3650")
            print("\nFor more help:")
            print("  precipgen simulate --help")
            return 1
        
        # Load parameters
        print(f"Loading parameters from {params_path}...")
        
        try:
            params = load_parameter_manifest(params_path)
        except DataLoadError as e:
            print(f"Error loading parameters: {e}")
            print("\nTroubleshooting:")
            print("  • Verify the parameter file path is correct")
            print("  • Ensure the file is valid JSON format")
            print("  • Check that analysis completed successfully")
            print("  • Re-run analysis if needed:")
            print(f"    precipgen analyze --station {args.station or 'STATION_ID'} --auto-detect")
            return 1
        
        station_id = params.get('station_id', args.station or 'unknown')
        print(f"Parameters loaded for station: {station_id}")
        
        # Handle status request
        if args.status:
            print(f"Simulation status for station {station_id}:")
            print(f"Parameter file: {params_path}")
            print(f"Data period: {params.get('data_period', {}).get('start', 'Unknown')} to {params.get('data_period', {}).get('end', 'Unknown')}")
            print(f"Trend mode: {'Enabled' if args.trends else 'Disabled'}")
            return 0
        
        # Handle interactive mode
        if args.interactive:
            interactive_simulation(params, args.seed)
            return 0
        
        # Validate simulation parameters
        if args.days <= 0:
            print("Error: Number of days must be positive")
            print("Example: --days 3650 (for ~10 years)")
            return 1
        
        if args.days > 50000:
            print(f"Warning: Simulating {args.days} days ({args.days/365.25:.1f} years) - this may take a while")
            print("Consider using a smaller number of days for faster results")
            print("Typical values: 3650 days (~10 years), 7300 days (~20 years)")
        
        # Generate synthetic precipitation
        print(f"Generating {args.days} days of synthetic precipitation...")
        if args.seed is not None:
            print(f"Using random seed: {args.seed}")
        
        try:
            precip_data = generate_synthetic_precipitation(
                params, args.days, args.seed, args.trends, args.start_year
            )
        except Exception as e:
            logger.error(f"Simulation failed: {e}")
            print(f"Error during simulation: {e}")
            return 1
        
        print(f"✓ Generated {len(precip_data)} days of synthetic data")
        
        # Calculate statistics
        print("Calculating simulation statistics...")
        stats = calculate_simulation_statistics(precip_data)
        
        # Create metadata
        metadata = {
            'station_id': station_id,
            'parameter_file': params_path,
            'simulation_period': {
                'start': str(precip_data.index.min().date()),
                'end': str(precip_data.index.max().date()),
                'total_days': len(precip_data)
            },
            'simulation_config': {
                'random_seed': args.seed,
                'trend_mode': args.trends,
                'start_year': args.start_year if args.trends else None
            },
            'generated_at': datetime.now().isoformat()
        }
        
        # Save output
        print("Saving simulation output...")
        saved_files = save_simulation_output(
            precip_data, stats, metadata, station_id, 
            args.format, args.output_file, dir_manager
        )
        
        for file_path in saved_files:
            print(f"✓ Saved: {file_path}")
        
        # Summary statistics
        total_precip = precip_data.sum()
        wet_days = (precip_data > 0.1).sum()
        wet_frequency = wet_days / len(precip_data)
        
        print(f"\nSimulation summary:")
        print(f"  Period: {precip_data.index.min().strftime('%Y-%m-%d')} to {precip_data.index.max().strftime('%Y-%m-%d')}")
        print(f"  Total precipitation: {total_precip:.1f}mm")
        print(f"  Wet days: {wet_days} ({wet_frequency:.1%})")
        print(f"  Mean daily precipitation: {precip_data.mean():.2f}mm")
        print(f"  Maximum daily precipitation: {precip_data.max():.1f}mm")
        
        if args.trends:
            print(f"  Trend mode: Enabled (start year: {args.start_year})")
        
        print(f"\nNext steps:")
        print(f"  1. Validate results: precipgen validate --station {station_id}")
        print(f"  2. View outputs: precipgen list-outputs --station {station_id}")
        
        return 0
        
    except Exception as e:
        logger.error(f"Unexpected error in simulate command: {e}")
        print(f"Error: {e}")
        return 1