#!/usr/bin/env python3
import argparse
import json
import logging
import os
import signal
import sys
import threading
import time
import urllib3
from concurrent.futures import ThreadPoolExecutor, as_completed

import boto3
from botocore.exceptions import ClientError, BotoCoreError

CAT_ART = """\
            ::
           :  :
         .'    '.
        /        \\
       |          |
       |  EYE  EYE|
       |          |
       \\   (__)   /
        \\        /
         '::::::'
""".replace('EYE', '\033[38;2;255;215;0m()__()\033[0m')

CAT_ART = """\
      ,_     _
      |\\\\_,-~/
      / _  _ |    ,--.
     (  X  X )   / ,-'
      \\  _T_/-._( (
      /         `. \\
     |         _  \\ |
      \\ \\ ,  /      |
       || |-_\\__   /
      ((_/`(____,-'
""".replace('X', '\033[38;2;255;215;0mO\033[0m')

logger = logging.getLogger(__name__)

class ColorFormatter(logging.Formatter):
    COLORS = {
        logging.DEBUG: '\033[94m',      # Blue
        logging.INFO: '\033[0m',        # Default
        logging.WARNING: '\033[93m',    # Yellow
        logging.ERROR: '\033[91m',      # Red
        logging.CRITICAL: '\033[91m\033[1m', # Bold Red
    }
    RESET = '\033[0m'

    def format(self, record):
        color = self.COLORS.get(record.levelno, self.RESET)
        log_fmt = f'%(asctime)s - {color}%(levelname)s{self.RESET} - {color}%(message)s{self.RESET}'
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)

# Global flag to signal threads to stop when we need to gracefully shutdown (e.g., token expiration or user interrupt)
shutdown_event = threading.Event()

# Thread-local storage for boto3 clients, as they are not thread-safe
tls = threading.local()

def get_ebs_client(region, verify_ssl=True):
    """Returns a thread-local boto3 EBS client."""
    if not hasattr(tls, 'client'):
        # Creating a fresh session for each thread
        session = boto3.Session(region_name=region) if region else boto3.Session()
        tls.client = session.client('ebs', verify=verify_ssl)
    return tls.client

def load_state(state_file, output_file, block_size):
    """Loads the completed blocks from the state file and validates they can exist on disk."""
    if os.path.exists(state_file):
        try:
            with open(state_file, 'r') as f:
                data = json.load(f)
                completed = set(data.get('completed_blocks', []))
                
            if not completed:
                return set()
                
            if not os.path.exists(output_file):
                logger.warning(f"Output file '{output_file}' does not exist! "
                               f"Discarding state file to prevent creating corrupted/empty image.")
                return set()
                
            # Basic validation that file is large enough to hold these blocks
            file_size = os.path.getsize(output_file)
            valid = set()
            for b in completed:
                if (b + 1) * block_size <= file_size or file_size == 0:
                     # If file_size is 0, it means it might be preallocated later, but since it exists, we'll keep it.
                     # Actually, if it's smaller, we discard the block.
                     valid.add(b)
                
            if len(valid) < len(completed):
                logger.warning(f"Output file truncated! Invalidated {len(completed) - len(valid)} blocks.")
                
            logger.info(f"Loaded state file. {len(valid)} valid blocks previously completed.")
            return valid
        except Exception as e:
            logger.warning(f"Failed to load state file: {e}. Starting fresh.")
    return set()

def save_state(state_file, completed_blocks):
    """Atomically saves the state file."""
    tmp_file = state_file + '.tmp'
    try:
        with open(tmp_file, 'w') as f:
            json.dump({'completed_blocks': list(completed_blocks)}, f)
        os.replace(tmp_file, state_file)
    except Exception as e:
        logger.error(f"Failed to save state file safely: {e}")

def signal_handler(signum, frame):
    """Handles Ctrl+C (SIGINT)."""
    logger.info("\nInterrupt received. Shutting down gracefully. Please wait for active downloads to finish and state to save...")
    shutdown_event.set()

def format_bytes(size):
    power = 2**10
    n = 0
    power_labels = {0 : '', 1: 'K', 2: 'M', 3: 'G', 4: 'T'}
    while size > power:
        size /= power
        n += 1
    return f"{size:.2f} {power_labels[n]}B"

def download_block_worker(snapshot_id, block_index, block_token, block_size, region, verify_ssl):
    """Worker function to download a single EBS snapshot block."""
    if shutdown_event.is_set():
        return False, block_index, "Shutdown requested"
    
    client = get_ebs_client(region, verify_ssl)
    try:
        response = client.get_snapshot_block(
            SnapshotId=snapshot_id,
            BlockIndex=block_index,
            BlockToken=block_token
        )
        # BlockData is a StreamingBody, read the bytes
        block_data = response['BlockData'].read()
        return True, block_index, block_data

    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        # If tokens expire during run, we want to catch it and cleanly shut down
        if error_code in ['ExpiredTokenException', 'UnrecognizedClientException', 'InvalidAccessKeyId', 'AccessDenied']:
            print("") # Drop to a new line so we don't print over the \r progress bar
            logger.error(f"AWS Credentials rejected/expired in thread (Error Code: {error_code}). Triggering graceful shutdown.")
            shutdown_event.set()
        return False, block_index, str(e)
    except Exception as e:
        return False, block_index, str(e)

def main():
    parser = argparse.ArgumentParser(description="Download an AWS EBS Snapshot to a local raw disk image.")
    parser.add_argument("--snapshot-id", required=True, help="AWS EBS Snapshot ID (e.g., snap-0123456789abcdef0)")
    parser.add_argument("--output-file", required=True, help="Path to the local output raw disk image")
    parser.add_argument("--region", help="AWS Region (optional, uses default from environment if omitted)")
    parser.add_argument("--state-file", help="Path to state file for resumption (defaults to <output-file>.state.json)")
    parser.add_argument("--workers", type=int, default=16, help="Number of concurrent download threads (default: 16)")
    parser.add_argument("--log-file", help="Path to write all log output (defaults to <output-file>.log)")
    parser.add_argument("--no-verify-ssl", action="store_true", help="Disable SSL certificate verification for AWS API calls")
    
    args = parser.parse_args()
    
    snapshot_id = args.snapshot_id
    output_file = args.output_file
    region = args.region
    state_file = args.state_file if args.state_file else f"{output_file}.state.json"
    log_file = args.log_file if args.log_file else f"{output_file}.log"
    workers = args.workers
    verify_ssl = not args.no_verify_ssl

    if not verify_ssl:
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    
    # Send all logs to the file
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(log_formatter)
    
    # Send all logs to the console as well
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(ColorFormatter())
    
    root_logger = logging.getLogger()
    root_logger.handlers.clear()
    root_logger.addHandler(file_handler)
    root_logger.addHandler(console_handler)
    root_logger.setLevel(logging.INFO)
    
    # Suppress verbose AWS botocore logs ("Found credentials...") so it doesn't break the progress bar visual
    logging.getLogger('botocore').setLevel(logging.WARNING)
    logging.getLogger('boto3').setLevel(logging.WARNING)
    logging.getLogger('urllib3').setLevel(logging.WARNING)

    print(CAT_ART)
    logger.info(f"Initializing bsnap...")
    logger.info(f"Detailed logs will also be sent to: {log_file}")

    # Register interrupt handler
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    logger.info(f"Connecting to AWS to initialize download for snapshot {snapshot_id}...")
    
    # Initialize main thread client
    try:
        main_session = boto3.Session(region_name=region) if region else boto3.Session()
        main_client = main_session.client('ebs', verify=verify_ssl)
    except BotoCoreError as e:
        logger.error(f"Failed to initialize AWS client: {e}")
        sys.exit(1)

    # Initial pagination call to get volume and block size details
    # We do a trial run on list_snapshot_blocks just to check connectivity and sizing
    try:
        logger.info("Fetching snapshot metadata...")
        initial_resp = main_client.list_snapshot_blocks(SnapshotId=snapshot_id, MaxResults=100)
        volume_size_gb = initial_resp.get('VolumeSize', 0)
        block_size = initial_resp.get('BlockSize', 524288) # Default 512 KiB
        # volume_size in list_snapshot_blocks is typically in GiB
        total_volume_bytes = volume_size_gb * 1024 * 1024 * 1024
        logger.info(f"Volume Size: {volume_size_gb} GiB ({format_bytes(total_volume_bytes)})")
        logger.info(f"Block Size: {block_size} bytes")
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        if error_code in ['ExpiredTokenException', 'UnrecognizedClientException', 'InvalidAccessKeyId']:
            logger.error("AWS credentials are invalid or expired before starting.")
            logger.error("Please refresh your AWS credentials and try again.")
        else:
            logger.error(f"Failed to list snapshot blocks: {e}")
        sys.exit(1)

    completed_blocks = load_state(state_file, output_file, block_size)
    
    # Pre-allocate or open output file
    # We open in r+b if it exists so we can randomly write blocks, or w+b if it doesn't
    mode = 'r+b' if os.path.exists(output_file) else 'w+b'
    
    logger.info(f"Opening local image file: {output_file}")
    try:
        f_out = open(output_file, mode)
        # Ensure sparse file size
        if total_volume_bytes > 0:
            f_out.seek(total_volume_bytes - 1)
            f_out.write(b'\0')
    except IOError as e:
        logger.error(f"Failed to open/allocate output file: {e}")
        sys.exit(1)

    file_lock = threading.Lock()
    
    logger.info("Starting block download.")
    total_blocks_visited = 0
    downloaded_this_session = 0
    start_time = time.time()
    last_progress_time = start_time
    
    next_token = None
    
    try:
        # Manual pagination over the snapshot blocks since botocore doesn't have a built-in paginator for this
        while not shutdown_event.is_set():
            kwargs = {'SnapshotId': snapshot_id, 'MaxResults': 10000}
            if next_token:
                kwargs['NextToken'] = next_token
                
            page = main_client.list_snapshot_blocks(**kwargs)
            
            blocks = page.get('Blocks', [])
            
            # Filter blocks that are already completed
            blocks_to_download = [b for b in blocks if b['BlockIndex'] not in completed_blocks]
            
            total_blocks_visited += len(blocks)
            
            if not blocks_to_download:
                continue

            with ThreadPoolExecutor(max_workers=workers) as executor:
                # Submit tasks for this page
                futures = {
                    executor.submit(
                        download_block_worker, 
                        snapshot_id, 
                        b['BlockIndex'], 
                        b['BlockToken'], 
                        block_size, 
                        region,
                        verify_ssl
                    ): b for b in blocks_to_download
                }
                
                for future in as_completed(futures):
                    success, b_idx, result_or_err = future.result()
                    
                    if success:
                        block_data = result_or_err
                        # Write to file and update stats inside the lock
                        with file_lock:
                            f_out.seek(b_idx * block_size)
                            f_out.write(block_data)
                        
                            completed_blocks.add(b_idx)
                            downloaded_this_session += 1
                            
                            current_time = time.time()
                            if current_time - last_progress_time >= 1.0:
                                last_progress_time = current_time
                                duration = current_time - start_time
                                downloaded_mb = (downloaded_this_session * block_size) / (1024 * 1024)
                                throughput = downloaded_mb / duration if duration > 0 else 0
                                raw_msg = (f"[+] Progress: {len(completed_blocks)} / {total_blocks_visited} blocks retrieved | "
                                           f"Data transferred: {downloaded_mb:.2f} MB | "
                                           f"Elapsed time: {duration:.1f}s | "
                                           f"Throughput: {throughput:.2f} MB/s")
                                
                                col_msg = (f"\r\033[K\033[96m[+]\033[0m Progress: \033[1m{len(completed_blocks)} / {total_blocks_visited}\033[0m blocks retrieved | "
                                           f"Data: \033[92m{downloaded_mb:.2f} MB\033[0m | "
                                           f"Time: \033[93m{duration:.1f}s\033[0m | "
                                           f"Speed: \033[94m{throughput:.2f} MB/s\033[0m")
                                
                                # \033[K cleanly erases the line without relying on fixed padding
                                sys.stdout.write(col_msg)
                                sys.stdout.flush()
                                
                                # Do NOT log the progress message via logger.info, to keep console clean.
                                # But we want it in the log file without colors:
                                file_handler.stream.write(raw_msg + '\n')
                                file_handler.stream.flush()
                            
                            if downloaded_this_session % 1000 == 0:
                                # Periodic state save to disk
                                save_state(state_file, completed_blocks)
                    else:
                        if shutdown_event.is_set():
                            # Don't spam warnings if we are actively shutting down
                            break
                        else:
                            print("")
                            logger.warning(f"Failed to download block {b_idx}: {result_or_err}")
                            # Let it retry gracefully by NOT adding it to completed_blocks
                            
            # Save state at the end of every page processing to ensure resiliency
            save_state(state_file, completed_blocks)
            
            next_token = page.get('NextToken')
            if not next_token or shutdown_event.is_set():
                break

    except ClientError as e:
        print("")
        logger.error(f"Error while enumerating blocks: {e}")
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        if error_code in ['ExpiredTokenException', 'UnrecognizedClientException', 'InvalidAccessKeyId']:
            logger.error("Tokens expired during volume iteration! State saved. Please refresh tokens and resume.")
            shutdown_event.set()

    except Exception as e:
        print("")
        logger.error(f"Unexpected error: {e}")
    finally:
        # Move the console to a new line before final summary
        print("")
        
        logger.info("Closing file and saving final state...")
        f_out.close()
        save_state(state_file, completed_blocks)
        
        duration = time.time() - start_time
        downloaded_mb = (downloaded_this_session * block_size) / (1024 * 1024)
        throughput = downloaded_mb / duration if duration > 0 else 0
        
        logger.info(f"Finished session.")
        logger.info(f"Total time: {duration:.2f} seconds")
        logger.info(f"Data transferred: {downloaded_mb:.2f} MB")
        logger.info(f"Throughput: {throughput:.2f} MB/s")
        
        if shutdown_event.is_set():
            logger.info("Tool exited early due to shutdown event (e.g. Expired Token or Ctrl+C).")
            logger.info(f"Progress so far: {len(completed_blocks)} / {total_blocks_visited} blocks retrieved.")
            logger.info(f"You can resume the download by running the exact same command. "
                        f"State is saved in '{state_file}'.")
            sys.exit(1)
        else:
            if len(completed_blocks) >= total_blocks_visited and total_blocks_visited >= 0:
                logger.info(f"Blocks successfully retrieved: {len(completed_blocks)} / {total_blocks_visited}")
                logger.info("\033[92mSUCCESS: All blocks have been successfully retrieved!\033[0m")
                logger.info(f"Time: {duration:.1f}s, Data: {downloaded_mb:.2f} MB")
            else:
                missing = total_blocks_visited - len(completed_blocks)
                logger.warning(f"INCOMPLETE: Finished paginating, but {missing} blocks were not downloaded.")
                logger.info("Please re-run the exact same command to retry downloading the missing blocks.")
                sys.exit(1)

if __name__ == '__main__':
    main()
