#!/usr/bin/env python3
"""
Migrate stale url# lookup keys for publications whose url was updated to
include /pub/ but whose url# DynamoDB key was never updated.

For each publication record (u#<did>#pub / pub#<rkey>):
  - The `url` field is the authoritative current URL (e.g. .../pub/<rkey>)
  - The `url#<url>` / info key is what the .well-known handler looks up
  - When updatePublication ran after the /pub/ fix it updated `url` but not
    the url# key, leaving a stale key under the old URL

This script:
  1. Scans all entityType=publication items
  2. Checks whether the correct url# key exists
  3. If not: writes the correct key and deletes the old one (old URL = url
     with /pub/<rkey> replaced by /<rkey>)

Usage:
    python migrate-publication-urls.py [--dry-run] [--profile <aws-profile>]
"""

import argparse
import boto3
from boto3.dynamodb.conditions import Attr

TABLE = 'blogs'


def old_url(url: str) -> str | None:
    """Derive the pre-/pub/ URL from a current URL, or None if not applicable."""
    # current: https://lemma.pub/<did>/pub/<rkey>
    # old:     https://lemma.pub/<did>/<rkey>
    parts = url.rstrip('/').rsplit('/', 3)
    # parts: ['https://lemma.pub', '<did>', 'pub', '<rkey>']
    if len(parts) == 4 and parts[2] == 'pub':
        return f"{parts[0]}/{parts[1]}/{parts[3]}"
    return None


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

    print(f"{'[DRY RUN] ' if dry_run else ''}Scanning publications in {TABLE}...\n")

    paginator = table.meta.client.get_paginator('scan')
    pages = paginator.paginate(
        TableName=TABLE,
        FilterExpression='entityType = :et',
        ExpressionAttributeValues={':et': 'publication'},
    )

    fixed = skipped = already_ok = errors = 0

    for page in pages:
        for item in page['Items']:
            did = item['PK'].replace('u#', '').replace('#pub', '')
            rkey = item['SK'].replace('pub#', '')
            url = item.get('url', '')
            at_uri = item.get('atUri', '')

            if not url or not at_uri:
                print(f"  SKIP  {did}/{rkey} — missing url or atUri")
                skipped += 1
                continue

            # Check if the correct url# key already exists
            correct_key = {'PK': f'url#{url}', 'SK': 'info'}
            existing = table.get_item(Key=correct_key).get('Item')

            if existing:
                already_ok += 1
                continue

            # Correct key is missing — find and migrate the old key
            old = old_url(url)
            if not old:
                print(f"  SKIP  {did}/{rkey} — can't derive old URL from {url!r}")
                skipped += 1
                continue

            old_key = {'PK': f'url#{old}', 'SK': 'info'}
            old_item = table.get_item(Key=old_key).get('Item')

            print(f"  FIX   {did}/{rkey}")
            print(f"        old url# key: {old}")
            print(f"        new url# key: {url}")

            if not dry_run:
                try:
                    table.put_item(Item={
                        'PK': f'url#{url}',
                        'SK': 'info',
                        'entityType': 'url-lookup',
                        'pubAtUri': at_uri,
                    })
                    if old_item:
                        table.delete_item(Key=old_key)
                    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)
