import boto3
import time
import os
from botocore.exceptions import ClientError
from datetime import datetime, timedelta

def get_all_tables(athena_client, database_name):
    """Get all tables from the specified database."""
    try:
        tables = []
        paginator = athena_client.get_paginator('list_table_metadata')
        
        for page in paginator.paginate(CatalogName='AwsDataCatalog', DatabaseName=database_name):
            for table in page['TableMetadataList']:
                tables.append(table['Name'])
                
        return tables
    except ClientError as e:
        print(f"Error getting tables: {e}")
        return []

def execute_query(athena_client, database, query, output_location):
    """Execute Athena query and return execution ID."""
    try:
        response = athena_client.start_query_execution(
            QueryString=query,
            QueryExecutionContext={
                'Database': database
            },
            ResultConfiguration={
                'OutputLocation': output_location
            },
            WorkGroup='inventory_workgroup'
        )
        return response['QueryExecutionId']
    except ClientError as e:
        print(f"Error executing query: {e}")
        return None

def get_query_results(athena_client, query_execution_id):
    """Wait for query to complete and return results."""
    while True:
        try:
            response = athena_client.get_query_execution(QueryExecutionId=query_execution_id)
            state = response['QueryExecution']['Status']['State']
            
            if state in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
                break
                
            time.sleep(5)  # Wait 5 seconds before checking again
            
        except ClientError as e:
            print(f"Error getting query execution status: {e}")
            return None
    
    if state == 'SUCCEEDED':
        try:
            results = athena_client.get_query_results(QueryExecutionId=query_execution_id)
            return results
        except ClientError as e:
            print(f"Error getting query results: {e}")
            return None
    else:
        print(f"Query failed with state: {state}")
        return None

def download_results(s3_client, output_location, local_file):
    """Download query results from S3 to local file."""
    try:
        # Remove 's3://' prefix and split bucket and key
        bucket, key = output_location.replace('s3://', '').split('/', 1)
        s3_client.download_file(bucket, key, local_file)
        return True
    except ClientError as e:
        print(f"Error downloading results: {e}")
        return False

def main():
    # Initialize AWS clients
    athena_client = boto3.client('athena', region_name='us-east-2')
    s3_client = boto3.client('s3', region_name='us-east-2')
    
    # Configuration
    DATABASE = 'inventory_db'
    OUTPUT_BUCKET = 'backup-inventory-temp'  # Replace with your bucket
    OUTPUT_PREFIX = 'large_objects_query/'
    OUTPUT_LOCATION = f's3://{OUTPUT_BUCKET}/{OUTPUT_PREFIX}'
    LOCAL_OUTPUT_DIR = 'query_results'
    
    # Create local output directory if it doesn't exist
    os.makedirs(LOCAL_OUTPUT_DIR, exist_ok=True)
    
    # Get all tables in the database
    tables = get_all_tables(athena_client, DATABASE)
    
    if not tables:
        print("No tables found in the database")
        return
    
    dt = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d-01-00")

    # Create UNION ALL query for all tables
    queries = []
    for table in tables:
        query = f"""
        SELECT 
            '{table}' as source_table,
            bucket,
            key,
            size,
            last_modified_date
        FROM \"{table}\"
        WHERE size >= 5368709120  -- 5GB in bytes
        AND dt = '{dt}'
        """
        queries.append(query)
    
    final_query = " UNION ALL ".join(queries)
    print(final_query)
    # Execute query
    query_execution_id = execute_query(athena_client, DATABASE, final_query, OUTPUT_LOCATION)
    
    if not query_execution_id:
        print("Failed to execute query")
        return
    
    # Get results
    results = get_query_results(athena_client, query_execution_id)
    
    if not results:
        print("No results obtained from query")
        return
    
    # Get the output location from the query execution
    query_execution = athena_client.get_query_execution(QueryExecutionId=query_execution_id)
    output_location = query_execution['QueryExecution']['ResultConfiguration']['OutputLocation']
    
    # Download results
    local_file = os.path.join(LOCAL_OUTPUT_DIR, f'large_objects_{query_execution_id}.csv')
    if download_results(s3_client, output_location, local_file):
        print(f"Results downloaded to: {local_file}")
    else:
        print("Failed to download results")

if __name__ == "__main__":
    main()