#!/usr/bin/env python3
"""
Delete all records for a given DID: posts (published and draft),
publications, and subscriber records.

Usage:
    python delete-user-posts.py <did> [--dry-run] [--profile <aws-profile>]
"""

import argparse
import boto3

TABLE = 'blogs'
BATCH_SIZE = 25


def delete_partition(table, pk: str, dry_run: bool) -> tuple[int, int]:
    last_key = None
    deleted = errors = 0

    while True:
        kwargs = {
            'KeyConditionExpression': 'PK = :pk',
            'ExpressionAttributeValues': {':pk': pk},
            'ProjectionExpression': 'PK, SK',
        }
        if last_key:
            kwargs['ExclusiveStartKey'] = last_key

        result = table.query(**kwargs)
        items = result.get('Items', [])

        for i in range(0, len(items), BATCH_SIZE):
            batch = items[i:i + BATCH_SIZE]
            if not dry_run:
                try:
                    with table.batch_writer() as writer:
                        for item in batch:
                            writer.delete_item(Key={'PK': item['PK'], 'SK': item['SK']})
                    deleted += len(batch)
                except Exception as e:
                    print(f"  ERROR: {e}")
                    errors += len(batch)
            else:
                deleted += len(batch)

        last_key = result.get('LastEvaluatedKey')
        if not last_key:
            break

    return deleted, errors


def run(did: str, dry_run: bool, profile: str):
    session = boto3.Session(profile_name=profile)
    ddb = session.resource('dynamodb')
    table = ddb.Table(TABLE)

    prefix = '[DRY RUN] ' if dry_run else ''
    partitions = [
        (f'u#{did}#post',        'posts'),
        (f'u#{did}#pub',         'publications'),
        (f'u#{did}#subscribers', 'subscriber records'),
    ]

    total_deleted = total_errors = 0
    for pk, label in partitions:
        print(f"{prefix}Deleting {label} ({pk})...")
        deleted, errors = delete_partition(table, pk, dry_run)
        print(f"  {'Would delete' if dry_run else 'Deleted'}: {deleted}  Errors: {errors}\n")
        total_deleted += deleted
        total_errors += errors

    print(f"{prefix}Done. Total {'would delete' if dry_run else 'deleted'}: {total_deleted}")
    if total_errors:
        print(f"  Total errors: {total_errors}")


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('did', help="DID of the user whose records to delete")
    ap.add_argument('--dry-run', action='store_true', help="Print what would change without writing")
    ap.add_argument('--profile', default='production', help="AWS profile (default: production)")
    args = ap.parse_args()
    run(did=args.did, dry_run=args.dry_run, profile=args.profile)
