#!/usr/bin/env python3
"""
Restore GSI1SK on all published blog-post items to the correct
`${publishedAt}#${rkey}` format, clamping any future publishedAt to now.

Items that already have the correct format are skipped.

Usage:
    python migrate-feed-sort-key.py [--dry-run] [--profile <aws-profile>]
"""

import argparse
import boto3
from datetime import datetime, timezone

TABLE = 'blogs'


def correct_gsi1sk(published_at: str, rkey: str, now: str) -> str:
    clamped = published_at if published_at <= now else now
    return f'{clamped}#{rkey}'


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

    now = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.') + \
          f'{datetime.now(timezone.utc).microsecond // 1000:03d}Z'

    print(f"{'[DRY RUN] ' if dry_run else ''}Restoring GSI1SK on blog-post#published items in {TABLE}...\n")

    paginator = table.meta.client.get_paginator('query')
    pages = paginator.paginate(
        TableName=TABLE,
        IndexName='FeedIndex',
        KeyConditionExpression='GSI1PK = :pk',
        ExpressionAttributeValues={':pk': 'blog-post#published'},
    )

    fixed = skipped = already_ok = errors = 0

    for page in pages:
        for item in page['Items']:
            pk = item['PK']
            sk = item['SK']
            gsi1sk = item.get('GSI1SK', '')
            published_at = item.get('publishedAt', '')

            if not sk.startswith('p#'):
                print(f"  SKIP  {pk} / {sk} — unexpected SK format")
                skipped += 1
                continue

            if not published_at:
                print(f"  SKIP  {pk} / {sk} — missing publishedAt")
                skipped += 1
                continue

            rkey = sk[len('p#'):]
            expected = correct_gsi1sk(published_at, rkey, now)

            if gsi1sk == expected:
                already_ok += 1
                continue

            if not dry_run:
                print(f"  FIX   {pk} / {sk}")
                print(f"        GSI1SK: {gsi1sk!r}  →  {expected!r}")
                try:
                    table.update_item(
                        Key={'PK': pk, 'SK': sk},
                        UpdateExpression='SET GSI1SK = :new_sk',
                        ExpressionAttributeValues={':new_sk': expected},
                    )
                    fixed += 1
                except Exception as e:
                    print(f"        ERROR: {e}")
                    errors += 1
            else:
                fixed += 1

    print(f"\n{'[DRY RUN] ' if dry_run else ''}Done.")
    print(f"  Fixed:      {fixed}")
    print(f"  Already ok: {already_ok}")
    print(f"  Skipped:    {skipped}")
    if errors:
        print(f"  Errors:     {errors}")


if __name__ == '__main__':
    ap = argparse.ArgumentParser(description=__doc__)
    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(dry_run=args.dry_run, profile=args.profile)
