import argparse
import boto3
from collections import defaultdict
import os

parser = argparse.ArgumentParser(description='')
parser.add_argument("--s3-bucket-name", type=str, help="S3 bucket name")
parser.add_argument("--upload-index-files", action="store_true", help="Upload index files to S3 bucket")
args = parser.parse_args()

if args.s3_bucket_name is None:
    print("Error: --s3-bucket-name is required")
    exit(1)

s3 = boto3.resource("s3")
bucket = s3.Bucket(args.s3_bucket_name)

if args.upload_index_files:
    print("Uploading index files to S3 bucket")
    # List all files in dist directory, finding all index.html files
    for root, dirs, files in os.walk("dist"):
        for file in files:
            if file == "index.html":
                file_path = os.path.join(root, file)
                # Remove the 'dist/' prefix from the S3 key
                s3_key = file_path[5:] if file_path.startswith("dist/") else file_path
                print(f" * Uploading {file_path} to s3://{args.s3_bucket_name}/{s3_key}")
                bucket.upload_file(
                    file_path,
                    s3_key,
                    ExtraArgs={'ContentType': 'text/html'}
                )
    exit(0)

# Create a dictionary to store files by directory
directories = defaultdict(list)

print(f"Listing objects in bucket {args.s3_bucket_name}")
for obj in bucket.objects.all():
    print(f" - Processing {obj.key}")
    
    # Get all parent directories of this object
    parts = obj.key.split('/')
    for i in range(len(parts)):
        dir_path = '/'.join(parts[:i])
        if dir_path:  # Skip empty string for root
            directories[dir_path].append(obj.key)
    # Add to root directory
    directories[''].append(obj.key)

# Create index.html for each directory
os.makedirs('dist', exist_ok=True)

def create_index_file(dir_path, files):
    # Create subdirectories in dist if needed
    full_path = os.path.join('dist', dir_path) if dir_path else 'dist'
    os.makedirs(full_path, exist_ok=True)
    
    file_to_write = os.path.join(full_path, 'index.html')
    print(f"Writing index.html to {file_to_write}")
    with open(file_to_write, 'w') as f:
        f.write("<!DOCTYPE html>\n<html>\n<head>\n")
        f.write(f"<title>Index of {dir_path if dir_path else '/'}</title>\n")
        f.write("</head>\n<body>\n")
        f.write(f"<h1>Index of {dir_path if dir_path else '/'}</h1>\n")
        
        # Add parent directory link if not in root
        if dir_path:
            parent = '/'.join(dir_path.split('/')[:-1])
            f.write(f'<a href="../index.html">../</a><br>\n')
        
        # Track unique subdirectories and files
        subdirs = set()
        regular_files = set()
        
        # List all files in this directory
        current_prefix = dir_path + '/' if dir_path else ''
        for file_path in sorted(files):
            if file_path.startswith(current_prefix):
                relative_path = file_path[len(current_prefix):]
                if '/' in relative_path:
                    # This is a subdirectory
                    subdir = relative_path.split('/')[0]
                    subdirs.add(subdir)
                else:
                    # This is a file
                    regular_files.add(relative_path)
        
        # Write subdirectories first
        for subdir in sorted(subdirs):
            f.write(f'<a href="{subdir}/index.html">{subdir}/</a><br>\n')
            
        # Write regular files
        for file_name in sorted(regular_files):
            if file_name != 'index.html':  # Skip the index.html file itself
                f.write(f'<a href="/{current_prefix}{file_name}">{file_name}</a><br>\n')
        
        f.write("</body>\n</html>\n")

# Generate index files for all directories
for dir_path, files in directories.items():
    print(f"Generating index.html for directory: {dir_path if dir_path else 'root'}")
    create_index_file(dir_path, files)
