# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "boto3",
#     "click",
#     "duckdb",
#     "psutil",
# ]
# ///

import duckdb
import json
import sys
from typing import Dict, List
import boto3
import click
from datetime import datetime
import psutil

conn = duckdb.connect(':memory:')
s3_client = boto3.client('s3')


inventory_buckets = ["pde-inventory-us-east-1", "pde-inventory-us-east-2", "pde-inventory-us-west-2"]

def parse_s3_path(s3_path: str) -> Dict[str, str]:
    """Parse an S3 path into bucket and key components"""
    path = s3_path.replace('s3://', '')
    bucket = path.split('/')[0]
    key = '/'.join(path.split('/')[1:])
    return {'bucket': bucket, 'key': key}

def get_file_paths_from_manifest(manifest_path: str, inventory_bucket: str) -> Dict:
    """Read and parse the manifest file from S3"""
    s3_client = boto3.client('s3')
    s3_parts = parse_s3_path(manifest_path)
    response = s3_client.get_object(
        Bucket=s3_parts['bucket'],
        Key=s3_parts['key']
    )
    manifest_content = json.loads(response['Body'].read().decode('utf-8'))
    files = [f"s3://{inventory_bucket}/{file['key']}" for file in manifest_content['files']]
    return files

def create_inventory_view(conn: duckdb.DuckDBPyConnection, manifest_path: str, inventory_bucket: str) -> None:
    """Create a DuckDB view from the S3 inventory files"""
    files = get_file_paths_from_manifest(manifest_path, inventory_bucket)
        
    # Create a table with the appropriate schema
    create_view_sql = f"""
    CREATE VIEW s3_inventory AS
    SELECT * FROM read_parquet({files})
    """
    conn.execute(create_view_sql)
    

def query_inventory(conn: duckdb.DuckDBPyConnection, sql_query: str, output_file: str) -> List[Dict]:
    """Execute a SQL query against the inventory"""
    return conn.sql(sql_query).write_csv(output_file, header=False)

def find_inventory_bucket(bucket: str) -> str:
    """Find the inventory bucket for a given bucket name"""
    for b in inventory_buckets:
        try:
            # List objects with the bucket name as prefix and delimiter to check only root level
            response = s3_client.list_objects_v2(
                Bucket=b,
                Prefix=f"{bucket}/",
                Delimiter='/'
            )
            
            # If we get any results back, this inventory bucket contains our target bucket
            if response.get('CommonPrefixes') or response.get('Contents'):
                return b
        except Exception:
            # Skip this bucket if we can't access it or other errors occur
            continue
    raise Exception(f"No inventory found for {bucket}")

def get_inventory_partition(inventory_bucket: str, bucket: str) -> str:
    """Get the inventory partition for a given bucket name"""
    try:
        # List objects with the bucket/20 prefix to find inventory dates
        response = s3_client.list_objects_v2(
            Bucket=inventory_bucket,
            Prefix=f"{bucket}/default_entire_bucket_daily_parquet/20",
            Delimiter='/'
        )

        # Get all the prefixes (folders) and sort them in descending order
        inventory_dates = []
        if 'CommonPrefixes' in response:
            for prefix in response['CommonPrefixes']:
                inventory_dates.append(prefix['Prefix'].rstrip('/'))
        
        # Sort in descending order to get latest first
        inventory_dates.sort(reverse=True)
        
        if inventory_dates:
            return inventory_dates[0]  # Return the latest inventory date
        
        return None

    except Exception as e:
        print(f"Error getting inventory partition: {str(e)}")
        return None

def check_available_ram():
    """Check if system has at least 8GB of RAM available"""
    MIN_RAM_GB = 8
    available_ram_gb = psutil.virtual_memory().available / (1024 ** 3)  # Convert bytes to GB
    
    if available_ram_gb < MIN_RAM_GB:
        print(f"Warning: This script recommends at least {MIN_RAM_GB}GB of available RAM.")
        print(f"Current available RAM: {available_ram_gb:.1f}GB")
        
        if not click.confirm('Continue anyway? This might cause the script to fail or your system to become unresponsive for large inventory files. This is not a joke, the computer may crash.'):
            sys.exit(1)

@click.command()
@click.option('-b', '--bucket', help='Bucket for the manifest', prompt=True)
@click.option('--prefix', help='Filter objects by key prefix')
@click.option('--date-before', type=click.DateTime(), help='Filter objects modified before this date (YYYY-MM-DD)')
@click.option('--date-after', type=click.DateTime(), help='Filter objects modified after this date (YYYY-MM-DD)')
@click.option('-o', '--output', help='Output file name (defaults to <bucket>-manifest.csv)')
def main(bucket: str, prefix: str, date_before: datetime, date_after: datetime, output: str):
    """Analyze S3 inventory data from a manifest file."""
    
    check_available_ram()
    
    # Set default output filenames if not provided
    base_output = output or f"{bucket}-manifest"
    large_output = f"{base_output}-large.csv"
    small_output = f"{base_output}-small.csv"
    
    inventory_bucket = find_inventory_bucket(bucket)
    print(f"Inventory bucket: {inventory_bucket}")
    inventory_partition = get_inventory_partition(inventory_bucket, bucket)
    print(f"Inventory partition: {inventory_partition}")
    manifest_path = f"s3://{inventory_bucket}/{inventory_partition}/manifest.json"
    print(f"Manifest path: {manifest_path}")
    
    try:
        # Create the view of inventory data
        create_inventory_view(conn, manifest_path, inventory_bucket)
        
        # Build the base query with filters
        base_query = ["""SELECT bucket, key, version_id FROM s3_inventory 
                     WHERE is_latest = true 
                     AND is_delete_marker = false"""]
        
        if prefix:
            base_query.append(f"AND key LIKE '{prefix}%'")
        if date_before:
            base_query.append(f"AND last_modified_date < '{date_before.isoformat()}'")
        if date_after:
            base_query.append(f"AND last_modified_date > '{date_after.isoformat()}'")
        
        # Create queries for large and small files
        large_query = " ".join(base_query + ["AND size >= 5368709120"])  # 5GB in bytes
        small_query = " ".join(base_query + ["AND size < 5368709120"])
        
        # Execute queries and write to separate files
        query_inventory(conn, large_query, large_output)
        query_inventory(conn, small_query, small_output)
        
        print(f"Large files (≥5GB) written to: {large_output}")
        print(f"Small files (<5GB) written to: {small_output}")
                
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()