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

import sys

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

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
) -> List[str]:
    """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 = 4
    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
):
    """
    Generates CSV manifests for large and small files in a bucket for batch operations.
    
    NB! DuckDB does not handle bucket region redirects.
    Set AWS_REGION env var to the region that has the invetnory bucket 
    to work past this limitation.
    """

    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()