#!/usr/bin/env python3
"""
Wasabi Bucket Utilization Report Generator

Author: Gowtham Kamireddi
Co-Author: Francisco Santos

Features:
- Multiple credential management methods:
  * Read from AWS credentials file (~/.aws/credentials)
  * Manual input during execution
  * Hardcode credentials (WARNING: Only use in secure environments)
- Reusable sessions (boto3, requests)
- Exponential backoff retry logic
- Comprehensive error handling
- Progress indication with tqdm
- Proper logging
- Input validation

IMPORTANT: User must have WasabiAccountStatsAccess or WasabiBucketStatsAccess permissions.

Usage:
    python3 wasabi_utilization_report.py
    - Choose credential method when prompted
"""

import boto3
import requests
import csv
import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional, Tuple, List
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from tqdm import tqdm

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

# Custom exception for permission denied errors
class PermissionDeniedError(Exception):
    """Raised when 403 Forbidden error occurs (non-root credentials)"""
    pass

# Constants
S3_ENDPOINT = "https://s3.wasabisys.com"
STATS_BUCKET_URL = "https://stats.wasabisys.com/v1/standalone/utilizations/bucket"
MAX_RETRIES = 3
BACKOFF_FACTOR = 0.3
TIMEOUT_SECONDS = 20

# UI Constants
SEPARATOR_FULL = "=" * 70
SEPARATOR_ERROR = "❌" * 35
BUCKET_COLUMN_WIDTH = 45
ACTIVE_TB_WIDTH = 8
COLUMN_HEADERS = ["Bucket Name", "Active TB", "Deleted TB", "NumBillableObjects"]

# Permission messages
PERMISSION_MESSAGE = (
    "The Wasabi Stats API requires one of these permissions:\n"
    "   - WasabiAccountStatsAccess (account-level access)\n"
    "   - WasabiBucketStatsAccess (bucket-level access)"
)
ASSIGN_PERMISSIONS_MESSAGE = (
    "Please assign the required permission policies to your user\n"
    "and try again."
)


def bytes_to_tb(byte_value: int) -> float:
    """Convert bytes to terabytes."""
    return byte_value / (1024 ** 4)


def print_permission_denied_error(bucket: Optional[str] = None) -> None:
    """Print formatted permission denied error message."""
    print(f"\n{SEPARATOR_ERROR}")
    print("\n❌ PERMISSION DENIED - Missing Required Permissions")
    print(SEPARATOR_FULL)
    print(f"⚠️  {PERMISSION_MESSAGE}")
    if bucket:
        print(f"\nError at bucket: {bucket}")
    print(f"\n{ASSIGN_PERMISSIONS_MESSAGE}")
    print(SEPARATOR_FULL)
    print("\n❌ Please verify your permissions.\n")


class CredentialManager:
    """Manages secure credential retrieval from multiple sources."""
    
    AWS_CREDENTIALS_PATH = Path.home() / ".aws" / "credentials"
    
    @staticmethod
    def list_aws_profiles() -> List[str]:
        """
        Extract all profile names from AWS credentials file.
        
        Returns:
            List of profile names found in the credentials file
        """
        profiles = []
        
        if not CredentialManager.AWS_CREDENTIALS_PATH.exists():
            return profiles
        
        try:
            with open(CredentialManager.AWS_CREDENTIALS_PATH, 'r') as f:
                content = f.read()
            
            lines = content.split('\n')
            for line in lines:
                line = line.strip()
                
                # Check for profile header
                if line.startswith('[') and line.endswith(']'):
                    profile_name = line[1:-1].strip()
                    profiles.append(profile_name)
            
            return profiles
        except Exception as e:
            logger.warning(f"Failed to list AWS profiles: {e}")
            return profiles
    
    @staticmethod
    def parse_aws_credentials_file(profile: str = "default") -> Tuple[str, str]:
        """
        Read credentials from AWS credentials file.
        
        File format:
        [profile_name]
        aws_access_key_id = YOUR_ACCESS_KEY
        aws_secret_access_key = YOUR_SECRET_KEY
        """
        if not CredentialManager.AWS_CREDENTIALS_PATH.exists():
            raise FileNotFoundError(
                f"AWS credentials file not found at {CredentialManager.AWS_CREDENTIALS_PATH}"
            )
        
        try:
            with open(CredentialManager.AWS_CREDENTIALS_PATH, 'r') as f:
                content = f.read()
            
            lines = content.split('\n')
            current_profile = None
            access_key = None
            secret_key = None
            
            for line in lines:
                line = line.strip()
                
                # Check for profile header
                if line.startswith('[') and line.endswith(']'):
                    current_profile = line[1:-1].strip()
                    if current_profile == profile:
                        access_key = None
                        secret_key = None
                    continue
                
                # Parse key-value pairs
                if current_profile == profile and '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip()
                    
                    if key == 'aws_access_key_id':
                        access_key = value
                    elif key == 'aws_secret_access_key':
                        secret_key = value
            
            if not access_key or not secret_key:
                raise ValueError(
                    f"Profile '{profile}' not found or missing credentials in "
                    f"{CredentialManager.AWS_CREDENTIALS_PATH}"
                )
            
            logger.info(f"Loaded credentials from AWS profile '{profile}'")
            return access_key, secret_key
            
        except Exception as e:
            raise ValueError(f"Failed to parse AWS credentials file: {e}")
    
    @staticmethod
    def get_credentials_from_aws_file() -> Tuple[str, str]:
        """Get credentials from AWS credentials file with profile selection."""
        print("\n📁 AWS Credentials File Method")
        print("-" * 60)
        
        if not CredentialManager.AWS_CREDENTIALS_PATH.exists():
            print(f"❌ AWS credentials file not found at:")
            print(f"   {CredentialManager.AWS_CREDENTIALS_PATH}")
            print("\n   Create one with format:")
            print("   [default]")
            print("   aws_access_key_id = YOUR_ACCESS_KEY")
            print("   aws_secret_access_key = YOUR_SECRET_KEY")
            raise FileNotFoundError("AWS credentials file not found")
        
        # List all available profiles
        profiles = CredentialManager.list_aws_profiles()
        
        if not profiles:
            print("❌ No profiles found in AWS credentials file")
            raise ValueError("No profiles found in AWS credentials file")
        
        # Display available profiles
        print(f"\n✅ Found {len(profiles)} profile(s):\n")
        for idx, profile in enumerate(profiles, 1):
            marker = "✓ " if profile == "default" else "  "
            print(f"  {idx}) {marker}{profile}")
        
        print()
        
        # Get user choice
        while True:
            try:
                choice = input(f"Select profile [1-{len(profiles)}] (default: 1): ").strip()
                
                # Default to first profile if user just presses Enter
                if not choice:
                    choice = "1"
                
                choice_idx = int(choice) - 1
                
                if 0 <= choice_idx < len(profiles):
                    selected_profile = profiles[choice_idx]
                    break
                else:
                    print(f"❌ Invalid choice. Please enter a number between 1 and {len(profiles)}")
            except ValueError:
                print(f"❌ Invalid input. Please enter a number between 1 and {len(profiles)}")
        
        try:
            access_key, secret_key = CredentialManager.parse_aws_credentials_file(selected_profile)
            print(f"\n✅ Credentials loaded from profile '{selected_profile}'")
            return access_key, secret_key
        except Exception as e:
            logger.error(f"Failed to load credentials: {e}")
            raise
    
    @staticmethod
    def get_credentials_manual_input() -> Tuple[str, str]:
        """Get credentials from user input."""
        print("\n⌨️  Manual Input Method")
        print("-" * 60)
        print("⚠️  Ensure your account has the required permissions:")
        print("   - WasabiAccountStatsAccess (full access)")
        print("   - WasabiBucketStatsAccess (sub-user limited access)\n")
        
        access_key = input("Enter Access Key ID: ").strip()
        secret_key = input("Enter Secret Key: ").strip()
        
        if not access_key or not secret_key:
            raise ValueError("Root Access and Secret Key are required.")
        
        # Confirm input
        print(f"\nAccess Key (last 8 chars): ...{access_key[-8:]}")
        confirm = input("Proceed with these credentials? [Y/n] (default: yes): ").strip().lower()
        
        # Default to yes if user just presses Enter
        if not confirm:
            confirm = "y"
        
        if confirm != 'y':
            raise ValueError("Credential input cancelled by user")
        
        logger.info("Credentials provided via manual input")
        return access_key, secret_key
    
    @staticmethod
    def get_credentials_hardcoded() -> Tuple[str, str]:
        """Get hardcoded credentials."""
        print("\n🔐 Hardcoded Credentials Method")
        print("-" * 60)
        print("⚠️  WARNING: This method stores credentials in source code")
        print("   ONLY use in secure, isolated environments!")
        print("⚠️  Ensure your account has the required permissions:")
        print("   - WasabiAccountStatsAccess (account-level access)")
        print("   - WasabiBucketStatsAccess (bucket-level access)\n")
        
        # ========== EDIT THESE VALUES ONLY ==========
        ACCESS_KEY = "YOUR-ACCESS-KEY-HERE"
        SECRET_KEY = "YOUR-SECRET-KEY-HERE"
        # ==========================================
        
        if ACCESS_KEY == "YOUR-ACCESS-KEY-HERE" or SECRET_KEY == "YOUR-SECRET-KEY-HERE":
            raise ValueError(
                "Hardcoded credentials not configured. "
                "Update ACCESS_KEY and SECRET_KEY in the code."
            )
        
        print(f"✅ Using hardcoded credentials")
        print(f"Access Key (last 8 chars): ...{ACCESS_KEY[-8:]}")
        logger.warning("Using hardcoded credentials - ensure this is in a secure environment")
        
        return ACCESS_KEY.strip(), SECRET_KEY.strip()
    
    @staticmethod
    def prompt_credential_method() -> Tuple[str, str]:
        """
        Prompt user to select credential retrieval method.
        
        Returns:
            Tuple of (access_key, secret_key)
        """
        print("\n" + "=" * 70)
        print("CREDENTIAL COLLECTION")
        print("=" * 70)
        print("\n⚠️  IMPORTANT: User must have required permissions:")
        print("   - WasabiAccountStatsAccess (account-level access)")
        print("   - WasabiBucketStatsAccess (bucket-level access)\n")
        print("Select credential method:")
        print("  1) Read from AWS credentials file (~/.aws/credentials)")
        print("  2) Enter credentials manually")
        print("  3) Use hardcoded credentials")
        print()
        
        choice = input("Enter choice [1-3] (default: 1): ").strip()
        
        # Default to first option if user just presses Enter
        if not choice:
            choice = "1"
        
        methods = {
            '1': ('AWS Credentials File', CredentialManager.get_credentials_from_aws_file),
            '2': ('Manual Input', CredentialManager.get_credentials_manual_input),
            '3': ('Hardcoded', CredentialManager.get_credentials_hardcoded),
        }
        
        if choice not in methods:
            raise ValueError("Invalid choice. Please enter 1, 2, or 3.")
        
        method_name, method_func = methods[choice]
        print(f"\n📌 Using {method_name} method...\n")
        
        try:
            return method_func()
        except Exception as e:
            logger.error(f"Failed to get credentials: {e}")
            raise


class WasabiClient:
    """Manages Wasabi API interactions with retry logic."""
    
    def __init__(self, access_key: str, secret_key: str):
        self.access_key = access_key
        self.secret_key = secret_key
        
        # Create boto3 session (reusable for all S3 operations)
        self.boto_session = boto3.Session(
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key
        )
        self.s3_client = self.boto_session.client(
            "s3",
            endpoint_url=S3_ENDPOINT,
            region_name="us-east-1"
        )
        
        # Create requests session with retry strategy
        self.http_session = self._create_http_session()
    
    def _create_http_session(self) -> requests.Session:
        """Create a requests session with exponential backoff retry strategy."""
        session = requests.Session()
        
        retry_strategy = Retry(
            total=MAX_RETRIES,
            backoff_factor=BACKOFF_FACTOR,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET"]
        )
        
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        session.mount("http://", adapter)
        
        return session
    
    def _encode_auth_header(self) -> str:
        """Build authorization header for Wasabi Stats API.
        
        Per Wasabi documentation, the Stats API expects plain text credentials
        in the format: ACCESS_KEY:SECRET_KEY (NOT Base64 encoded)
        Ref: https://docs.wasabi.com/apidocs/authentication-with-wasabi-stats-api
        """
        # Ensure no leading/trailing whitespace in credentials
        access_key_clean = self.access_key.strip()
        secret_key_clean = self.secret_key.strip()
        
        # Return plain text format, not Base64
        credentials = f"{access_key_clean}:{secret_key_clean}"
        
        return credentials
    
    def list_buckets(self) -> List[str]:
        """List all buckets in the account."""
        try:
            response = self.s3_client.list_buckets()
            buckets = [b["Name"] for b in response.get("Buckets", [])]
            logger.info(f"Found {len(buckets)} bucket(s)")
            return buckets
        except Exception as e:
            logger.error(f"Failed to list buckets: {e}")
            raise
    
    def get_bucket_utilization(self, bucket: str) -> Optional[Tuple[float, float, int]]:
        """
        Fetch bucket utilization data.
        
        Returns:
            Tuple of (active_tb, deleted_tb, num_billable_objects) or None if no data
        """
        auth_header = self._encode_auth_header()
        headers = {
            "Authorization": auth_header
        }
        params = {
            "latest": "true",
            "pageNum": 0,
            "pageSize": 1
        }
        url = f"{STATS_BUCKET_URL}/{bucket}"
        
        try:
            response = self.http_session.get(
                url,
                headers=headers,
                params=params,
                timeout=TIMEOUT_SECONDS
            )
            response.raise_for_status()
            
            data = response.json()
            records = data.get("Records", [])
            
            if not records:
                logger.warning(f"No utilization data for bucket '{bucket}'")
                return None
            
            record = records[0]
            active_tb = round(bytes_to_tb(record.get("RawStorageSizeBytes", 0)), 4)
            deleted_tb = round(bytes_to_tb(record.get("DeletedStorageSizeBytes", 0)), 4)
            num_billable_objects = record.get("NumBillableObjects", 0)
            
            return active_tb, deleted_tb, num_billable_objects
            
        except requests.exceptions.Timeout:
            logger.error(f"Timeout while fetching data for '{bucket}'")
            return None
        except requests.exceptions.ConnectionError:
            logger.error(f"Connection error while fetching data for '{bucket}'")
            return None
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                logger.error(f"Authentication failed for '{bucket}' - check credentials")
            elif e.response.status_code == 403:
                # Raise exception to stop processing immediately
                raise PermissionDeniedError(
                    f"Permission denied for '{bucket}' - "
                    "The Wasabi Stats API requires WasabiAccountStatsAccess or WasabiBucketStatsAccess permissions"
                )
            elif e.response.status_code == 404:
                logger.warning(f"Bucket '{bucket}' utilization data not found")
            else:
                logger.error(f"HTTP error for '{bucket}': {e.response.status_code}")
            return None
        except Exception as e:
            logger.error(f"Failed to fetch utilization for '{bucket}': {e}")
            return None


def validate_credentials(client: WasabiClient) -> Tuple[bool, Optional[str]]:
    """
    Validate credentials by attempting to list buckets.
    
    Returns:
        Tuple of (is_valid, error_type) where error_type can be:
        - None if valid
        - '403' if permission denied (missing required permissions)
        - 'other' for other validation errors
    """
    try:
        client.list_buckets()
        logger.info("Credentials validated successfully")
        return True, None
    except Exception as e:
        error_msg = str(e)
        
        # Check for permission denied errors (403 / missing permissions)
        if "403" in error_msg or "AccessDenied" in error_msg or "Forbidden" in error_msg:
            return False, '403'
        else:
            logger.error(f"Credential validation failed: {e}")
            return False, 'other'


def main():
    """Main execution function."""
    print("\n" + "=" * 70)
    print("WASABI BUCKET UTILIZATION REPORT GENERATOR")
    print("=" * 70)
    
    try:
        # Get credentials using interactive method selection
        access_key, secret_key = CredentialManager.prompt_credential_method()
        
        # Create client
        client = WasabiClient(access_key, secret_key)
        
        # Validate credentials
        is_valid, error_type = validate_credentials(client)
        
        if not is_valid:
            if error_type == '403':
                # 403 error - permission denied - STOP immediately
                print(f"\n{SEPARATOR_ERROR}")
                print("\n❌ PERMISSION DENIED - Missing Required Permissions")
                print(SEPARATOR_FULL)
                print(f"⚠️  {PERMISSION_MESSAGE}")
                print("\nYou provided:")
                print(f"  Access Key (last 8 chars): ...{access_key[-8:]}")
                print(f"\n{ASSIGN_PERMISSIONS_MESSAGE}")
                print(SEPARATOR_FULL)
                print("\n❌ Process stopped. Please verify your permissions and try again.\n")
                sys.exit(1)
            else:
                # Other validation error
                logger.error("Failed to validate credentials")
                sys.exit(1)
        
        print("\n" + "=" * 70)
        
        # Generate timestamps
        filename_ts = datetime.now().strftime("%Y-%m-%d_%H%M%S")
        
        # List buckets
        logger.info("Fetching bucket list...")
        buckets = client.list_buckets()
        
        if not buckets:
            logger.warning("No buckets found in account")
            return
        
        # Prepare output
        output_dir = Path.cwd()
        filename = output_dir / f"wasabi_bucket_utilization_{filename_ts}.csv"
        
        print(f"\nProcessing {len(buckets)} bucket(s)...\n")
        print("Bucket Name" + " " * 30 + "Active TB    Deleted TB    Objects")
        print("-" * 70)
        
        # Process buckets with progress bar
        results = []
        bucket = None  # Initialize to avoid NameError in exception handlers
        
        try:
            with open(filename, mode="w", newline="", encoding="utf-8") as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow([
                    "Bucket Name", "Active TB", "Deleted TB", 
                    "NumBillableObjects"
                ])
                
                for bucket in tqdm(buckets, desc="Fetching utilization", unit="bucket"):
                    try:
                        result = client.get_bucket_utilization(bucket)
                        
                        if result:
                            active_tb, deleted_tb, num_billable_objects = result
                            print(f"{bucket:<{BUCKET_COLUMN_WIDTH}} {active_tb:>{ACTIVE_TB_WIDTH}} TB    {deleted_tb:>{ACTIVE_TB_WIDTH}} TB    {num_billable_objects:>12,}")
                            writer.writerow([
                                bucket, active_tb, deleted_tb, 
                                num_billable_objects
                            ])
                            results.append((bucket, active_tb, deleted_tb, num_billable_objects))
                        else:
                            print(f"{bucket:<{BUCKET_COLUMN_WIDTH}} {'No data':<{ACTIVE_TB_WIDTH}} {'No data':<{ACTIVE_TB_WIDTH}}")
                            writer.writerow([bucket, "No data", "No data", "No data"])
                    
                    except Exception as e:
                        # Re-raise any exception to stop processing immediately
                        raise
        
        except PermissionDeniedError as e:
            # 403 error - stop processing immediately
            # Clean up the CSV file since processing failed
            if Path(filename).exists():
                Path(filename).unlink()
            
            print_permission_denied_error(bucket)
            sys.exit(1)
        
        except Exception as e:
            # Any other error - stop processing immediately
            # Clean up the CSV file since processing failed
            if Path(filename).exists():
                Path(filename).unlink()
            
            print(f"\n{SEPARATOR_ERROR}")
            print("\n❌ ERROR - Process Stopped")
            print(SEPARATOR_FULL)
            if bucket:
                print(f"Error at bucket: {bucket}")
            print(f"Error message: {str(e)}")
            print(SEPARATOR_FULL)
            print("\n❌ Process stopped due to error.\n")
            sys.exit(1)
        
        # Summary
        print(f"\n{SEPARATOR_FULL}")
        print("SUMMARY")
        print(SEPARATOR_FULL)
        print(f"Total buckets: {len(buckets)}")
        print(f"Successfully processed: {len(results)}")
        
        if results:
            total_active_tb = sum(r[1] for r in results)
            total_deleted_tb = sum(r[2] for r in results)
            total_objects = sum(r[3] for r in results)
            
            print(f"\nAggregate Statistics:")
            print(f"  Total Active Storage: {total_active_tb:.4f} TB")
            print(f"  Total Deleted Storage: {total_deleted_tb:.4f} TB")
            print(f"  Total Billable Objects: {total_objects:,}")
        
        print(f"\nCSV file saved: {filename}")
        print(f"{SEPARATOR_FULL}\n")
        
        logger.info("Report generation completed successfully")
        
    except ValueError as e:
        logger.error(f"Configuration error: {e}")
        sys.exit(1)
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
        sys.exit(130)
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()