#!/usr/bin/env python3
"""
GuardDuty S3 Malware Scan Calculator

Calculates exponential backoff delays, timeout scenarios, and cost estimates
for AWS GuardDuty Malware Protection for S3.
"""

from dataclasses import dataclass
from typing import List, Tuple


@dataclass
class BackoffConfig:
    """Configuration for exponential backoff polling."""
    initial_delay: int = 1  # seconds
    max_delay: int = 30  # seconds
    timeout: int = 300  # seconds
    backoff_rate: float = 2.0


@dataclass
class CostConfig:
    """GuardDuty S3 pricing configuration."""
    cost_per_gb: float = 0.09
    cost_per_1000_objects: float = 0.215
    free_tier_gb: float = 1.0
    free_tier_objects: int = 1000
    has_free_tier: bool = False


@dataclass
class ScanTimeEstimate:
    """Estimated scan time breakdown."""
    file_size_mb: float
    estimated_scan_time: float  # seconds
    polling_overhead: float  # seconds
    total_time: float  # seconds
    confidence: str  # "Low", "Medium", "High"


def calculate_backoff_sequence(config: BackoffConfig) -> List[Tuple[int, float, float]]:
    """
    Calculate the exponential backoff sequence.

    Args:
        config: BackoffConfig with timing parameters

    Returns:
        List of tuples: (iteration, delay_seconds, cumulative_time)
    """
    sequence = []
    cumulative_time = 0.0
    iteration = 0
    current_delay = config.initial_delay

    while cumulative_time < config.timeout:
        iteration += 1

        # Apply exponential backoff with cap
        delay = min(current_delay, config.max_delay)
        cumulative_time += delay

        sequence.append((iteration, delay, cumulative_time))

        # Prepare next delay (exponential increase)
        current_delay *= config.backoff_rate

    return sequence


def analyze_timeout_scenario(config: BackoffConfig) -> dict:
    """
    Analyze what happens with given timeout configuration.

    Args:
        config: BackoffConfig with timing parameters

    Returns:
        Dictionary with analysis results
    """
    sequence = calculate_backoff_sequence(config)

    total_iterations = len(sequence)
    total_time = sequence[-1][2] if sequence else 0

    # Calculate average delay
    total_delay = sum(item[1] for item in sequence)
    avg_delay = total_delay / total_iterations if total_iterations > 0 else 0

    return {
        "total_iterations": total_iterations,
        "total_time": total_time,
        "average_delay": avg_delay,
        "sequence": sequence
    }


def calculate_scan_cost(
        num_files: int,
        avg_file_size_mb: float,
        config: CostConfig
) -> dict:
    """
    Calculate GuardDuty S3 malware scanning costs.

    Args:
        num_files: Number of files to scan
        avg_file_size_mb: Average file size in MB
        config: CostConfig with pricing parameters

    Returns:
        Dictionary with cost breakdown
    """
    # Convert to GB
    total_gb = (num_files * avg_file_size_mb) / 1024

    # Apply free tier if applicable
    billable_gb = max(0, total_gb - config.free_tier_gb) if config.has_free_tier else total_gb
    billable_objects = max(0, num_files - config.free_tier_objects) if config.has_free_tier else num_files

    # Calculate costs
    data_cost = billable_gb * config.cost_per_gb
    objects_cost = (billable_objects / 1000) * config.cost_per_1000_objects
    total_cost = data_cost + objects_cost

    return {
        "num_files": num_files,
        "avg_file_size_mb": avg_file_size_mb,
        "total_gb": total_gb,
        "billable_gb": billable_gb,
        "billable_objects": billable_objects,
        "data_cost": data_cost,
        "objects_cost": objects_cost,
        "total_cost": total_cost,
        "cost_per_file": total_cost / num_files if num_files > 0 else 0
    }


def print_backoff_table(sequence: List[Tuple[int, float, float]]) -> None:
    """Print formatted backoff sequence table."""
    print("\n" + "=" * 60)
    print("EXPONENTIAL BACKOFF SEQUENCE")
    print("=" * 60)
    print(f"{'Iteration':<12} {'Delay (s)':<15} {'Cumulative (s)':<15}")
    print("-" * 60)

    for iteration, delay, cumulative in sequence:
        print(f"{iteration:<12} {delay:<15.1f} {cumulative:<15.1f}")


def print_timeout_analysis(analysis: dict, config: BackoffConfig) -> None:
    """Print timeout scenario analysis."""
    print("\n" + "=" * 60)
    print("TIMEOUT ANALYSIS")
    print("=" * 60)
    print(f"Configuration:")
    print(f"  Initial Delay:  {config.initial_delay} seconds")
    print(f"  Max Delay:      {config.max_delay} seconds")
    print(f"  Timeout:        {config.timeout} seconds")
    print(f"  Backoff Rate:   {config.backoff_rate}x")
    print(f"\nResults:")
    print(f"  Total Iterations: {analysis['total_iterations']}")
    print(f"  Total Time:       {analysis['total_time']:.1f} seconds")
    print(f"  Average Delay:    {analysis['average_delay']:.2f} seconds")


def print_cost_breakdown(cost_data: dict, config: CostConfig) -> None:
    """Print cost breakdown table."""
    print("\n" + "=" * 60)
    print("COST BREAKDOWN")
    print("=" * 60)
    print(f"Scenario:")
    print(f"  Number of Files:    {cost_data['num_files']:,}")
    print(f"  Avg File Size:      {cost_data['avg_file_size_mb']:.2f} MB")
    print(f"  Total Data:         {cost_data['total_gb']:.3f} GB")

    if config.has_free_tier:
        print(f"\nFree Tier Applied:")
        print(f"  Free Data:          {config.free_tier_gb} GB")
        print(f"  Free Objects:       {config.free_tier_objects:,}")
        print(f"  Billable Data:      {cost_data['billable_gb']:.3f} GB")
        print(f"  Billable Objects:   {cost_data['billable_objects']:,}")

    print(f"\nCost Breakdown:")
    print(f"  Data Scanned:       ${cost_data['data_cost']:.4f}")
    print(f"  Objects Evaluated:  ${cost_data['objects_cost']:.4f}")
    print(f"  Total Cost:         ${cost_data['total_cost']:.4f}")
    print(f"  Cost Per File:      ${cost_data['cost_per_file']:.6f}")


def estimate_scan_time(file_size_mb: float) -> ScanTimeEstimate:
    """
    Estimate GuardDuty scan time based on file size.

    Note: These are estimates based on typical GuardDuty performance.
    Actual scan times may vary based on:
    - File type and complexity
    - Current GuardDuty load
    - AWS region
    - File structure (compressed, encrypted, etc.)

    Args:
        file_size_mb: File size in megabytes

    Returns:
        ScanTimeEstimate with time breakdown
    """
    # Baseline scan times (rough estimates from AWS documentation and field reports)
    # Small files: ~5-15 seconds
    # Medium files (10-50 MB): ~15-60 seconds
    # Large files (50-500 MB): ~60-180 seconds
    # Very large files (500+ MB): ~180-300+ seconds

    if file_size_mb < 1:
        base_scan_time = 5 + (file_size_mb * 5)  # 5-10 seconds
        confidence = "Medium"
    elif file_size_mb < 10:
        base_scan_time = 10 + (file_size_mb * 2)  # 10-30 seconds
        confidence = "Medium"
    elif file_size_mb < 50:
        base_scan_time = 30 + (file_size_mb * 0.8)  # 30-70 seconds
        confidence = "Medium"
    elif file_size_mb < 100:
        base_scan_time = 60 + (file_size_mb * 0.6)  # 60-120 seconds
        confidence = "Low"
    elif file_size_mb < 500:
        base_scan_time = 90 + (file_size_mb * 0.4)  # 90-290 seconds
        confidence = "Low"
    else:
        base_scan_time = 180 + (file_size_mb * 0.3)  # 180+ seconds
        confidence = "Low"

    # Add typical polling overhead (average of 2-3 polls before completion)
    # Using exponential backoff: first poll at 1s, second at 2s, third at 4s
    # Average overhead: ~7-10 seconds
    polling_overhead = 8.0

    total_time = base_scan_time + polling_overhead

    return ScanTimeEstimate(
        file_size_mb=file_size_mb,
        estimated_scan_time=base_scan_time,
        polling_overhead=polling_overhead,
        total_time=total_time,
        confidence=confidence
    )


def print_scan_time_estimate(estimate: ScanTimeEstimate) -> None:
    """Print scan time estimate."""
    print("\n" + "=" * 60)
    print("SCAN TIME ESTIMATE")
    print("=" * 60)
    print(f"File Size:              {estimate.file_size_mb:.2f} MB")
    print(f"Estimated Scan Time:    {estimate.estimated_scan_time:.1f} seconds "
          f"({estimate.estimated_scan_time / 60:.1f} minutes)")
    print(f"Polling Overhead:       {estimate.polling_overhead:.1f} seconds")
    print(f"Total Expected Time:    {estimate.total_time:.1f} seconds "
          f"({estimate.total_time / 60:.1f} minutes)")
    print(f"Confidence Level:       {estimate.confidence}")
    print(f"\nNote: Actual times may vary based on file type, complexity,")
    print(f"      AWS region load, and GuardDuty service capacity.")


def calculate_batch_scan_time(
        num_files: int,
        avg_file_size_mb: float,
        concurrent_scans: int = 1
) -> dict:
    """
    Calculate time to scan a batch of files.

    Args:
        num_files: Number of files to scan
        avg_file_size_mb: Average file size in MB
        concurrent_scans: Number of concurrent scans (GuardDuty processes in parallel)

    Returns:
        Dictionary with batch timing information
    """
    single_scan = estimate_scan_time(avg_file_size_mb)

    # If processing serially
    total_serial_time = single_scan.total_time * num_files

    # If processing in parallel (GuardDuty supports concurrent scans)
    total_parallel_time = (single_scan.total_time * num_files) / concurrent_scans

    return {
        "num_files": num_files,
        "avg_file_size_mb": avg_file_size_mb,
        "time_per_file": single_scan.total_time,
        "concurrent_scans": concurrent_scans,
        "total_serial_time": total_serial_time,
        "total_parallel_time": total_parallel_time,
        "total_serial_hours": total_serial_time / 3600,
        "total_parallel_hours": total_parallel_time / 3600
    }


def print_batch_scan_time(batch_data: dict) -> None:
    """Print batch scan time analysis."""
    print("\n" + "=" * 60)
    print("BATCH SCAN TIME ANALYSIS")
    print("=" * 60)
    print(f"Number of Files:        {batch_data['num_files']:,}")
    print(f"Average File Size:      {batch_data['avg_file_size_mb']:.2f} MB")
    print(f"Time Per File:          {batch_data['time_per_file']:.1f} seconds")
    print(f"\nSerial Processing:")
    print(f"  Total Time:           {batch_data['total_serial_time']:.1f} seconds")
    print(f"                        {batch_data['total_serial_hours']:.2f} hours")
    print(f"\nParallel Processing ({batch_data['concurrent_scans']} concurrent):")
    print(f"  Total Time:           {batch_data['total_parallel_time']:.1f} seconds")
    print(f"                        {batch_data['total_parallel_hours']:.2f} hours")
    print(f"\nNote: GuardDuty processes scans in parallel automatically.")
    print(f"      Actual throughput depends on AWS service capacity.")


def get_user_input() -> dict:
    """Get user input for calculations with defaults."""
    print("\nEnter calculation parameters (press Enter for defaults):\n")

    # Backoff configuration
    print("--- Backoff Configuration ---")
    initial_delay = input("Initial delay in seconds [1]: ").strip()
    initial_delay = int(initial_delay) if initial_delay else 1

    max_delay = input("Max delay in seconds [30]: ").strip()
    max_delay = int(max_delay) if max_delay else 30

    timeout = input("Timeout in seconds [300]: ").strip()
    timeout = int(timeout) if timeout else 300

    # Cost configuration
    print("\n--- Cost Calculation ---")
    num_files = input("Number of files per month [10000]: ").strip()
    num_files = int(num_files) if num_files else 10000

    avg_file_size = input("Average file size in MB [5.0]: ").strip()
    avg_file_size = float(avg_file_size) if avg_file_size else 5.0

    has_free_tier = input("Apply free tier? (y/n) [n]: ").strip().lower()
    has_free_tier = has_free_tier == 'y'

    return {
        'backoff_config': BackoffConfig(
            initial_delay=initial_delay,
            max_delay=max_delay,
            timeout=timeout,
            backoff_rate=2.0
        ),
        'num_files': num_files,
        'avg_file_size': avg_file_size,
        'cost_config': CostConfig(has_free_tier=has_free_tier)
    }


def run_interactive_mode():
    """Run calculator in interactive mode with user input."""
    print("=" * 60)
    print("GuardDuty S3 Malware Scan Calculator - Interactive Mode")
    print("=" * 60)

    user_input = get_user_input()

    # Scan time estimate
    print("\n\n### SCAN TIME ESTIMATE (Single File) ###")
    scan_estimate = estimate_scan_time(user_input['avg_file_size'])
    print_scan_time_estimate(scan_estimate)

    # Batch scan time
    print("\n\n### BATCH SCAN TIME ESTIMATE ###")
    batch_time = calculate_batch_scan_time(
        num_files=user_input['num_files'],
        avg_file_size_mb=user_input['avg_file_size'],
        concurrent_scans=10  # Typical concurrent capacity
    )
    print_batch_scan_time(batch_time)

    # Backoff analysis
    print("\n\n### BACKOFF ANALYSIS ###")
    analysis = analyze_timeout_scenario(user_input['backoff_config'])
    print_timeout_analysis(analysis, user_input['backoff_config'])

    # Show abbreviated backoff table for interactive mode
    print("\nBackoff Sequence (first 15 iterations):")
    print(f"{'Iteration':<12} {'Delay (s)':<15} {'Cumulative (s)':<15}")
    print("-" * 60)
    for iteration, delay, cumulative in analysis['sequence'][:15]:
        print(f"{iteration:<12} {delay:<15.1f} {cumulative:<15.1f}")
    if len(analysis['sequence']) > 15:
        print(f"... ({len(analysis['sequence']) - 15} more iterations)")

    # Cost calculation
    print("\n\n### COST ESTIMATE ###")
    cost = calculate_scan_cost(
        num_files=user_input['num_files'],
        avg_file_size_mb=user_input['avg_file_size'],
        config=user_input['cost_config']
    )
    print_cost_breakdown(cost, user_input['cost_config'])


def main():
    """Main function with example calculations."""
    import sys

    print("=" * 60)
    print("GuardDuty S3 Malware Scan Calculator")
    print("=" * 60)

    # Check if interactive mode requested
    if len(sys.argv) > 1 and sys.argv[1] in ['-i', '--interactive']:
        run_interactive_mode()
        return

    print("\nRunning in example mode. Use -i or --interactive for custom inputs.\n")

    # Example 1: Default configuration
    print("\n### EXAMPLE 1: Default Configuration (Small Files) ###")
    default_config = BackoffConfig(
        initial_delay=1,
        max_delay=30,
        timeout=300,
        backoff_rate=2.0
    )
    analysis = analyze_timeout_scenario(default_config)
    print_timeout_analysis(analysis, default_config)
    print_backoff_table(analysis['sequence'][:10])  # Show first 10 iterations

    # Example 2: Large files configuration
    print("\n\n### EXAMPLE 2: Large Files Configuration ###")
    large_file_config = BackoffConfig(
        initial_delay=5,
        max_delay=60,
        timeout=600,
        backoff_rate=2.0
    )
    analysis_large = analyze_timeout_scenario(large_file_config)
    print_timeout_analysis(analysis_large, large_file_config)

    # Example 3: Batch processing time
    print("\n\n### EXAMPLE 3: Batch Scan Time Analysis ###")
    batch_time = calculate_batch_scan_time(
        num_files=1000,
        avg_file_size_mb=10.0,
        concurrent_scans=10
    )
    print_batch_scan_time(batch_time)

    # Example 4: Cost calculation - Monthly volume
    print("\n\n### EXAMPLE 4: Monthly Cost Estimate (No Free Tier) ###")
    cost_config = CostConfig(has_free_tier=False)
    monthly_cost = calculate_scan_cost(
        num_files=10000,
        avg_file_size_mb=5.0,
        config=cost_config
    )
    print_cost_breakdown(monthly_cost, cost_config)

    # Example 5: Cost with free tier (new accounts)
    print("\n\n### EXAMPLE 5: Monthly Cost Estimate (With Free Tier) ###")
    cost_config_free = CostConfig(has_free_tier=True)
    monthly_cost_free = calculate_scan_cost(
        num_files=10000,
        avg_file_size_mb=5.0,
        config=cost_config_free
    )
    print_cost_breakdown(monthly_cost_free, cost_config_free)

    # Example 6: Different file sizes comparison
    print("\n\n### EXAMPLE 6: Cost Comparison by File Size ###")
    print("\n" + "=" * 60)
    print("COST COMPARISON (1000 files/month, No Free Tier)")
    print("=" * 60)
    print(f"{'File Size':<15} {'Total GB':<12} {'Total Cost':<15} {'Per File':<12}")
    print("-" * 60)

    file_sizes = [1, 5, 10, 50, 100]  # MB
    for size in file_sizes:
        cost = calculate_scan_cost(1000, size, CostConfig(has_free_tier=False))
        print(f"{size:>5} MB        {cost['total_gb']:>8.3f} GB   "
              f"${cost['total_cost']:>10.4f}      ${cost['cost_per_file']:>8.6f}")


if __name__ == "__main__":
    main()

# Run interactively
# python guardduty_s3_scan_estimates.py -i