#!/usr/bin/env python3
"""Migrate LocalBackend objects to the flat global content-addressed layout.

Old layouts (both broken):
    <root>/<repo_id>/<object_id>     — legacy per-repo paths
    <root>/objects/<object_id>       — double-objects bug (root was set to
                                       /data/musehub/objects but code appended
                                       "objects/" again)

Correct layout:
    <root>/<object_id>               — flat, globally content-addressed
                                       (root = musehub_objects_dir, default
                                        /data/musehub/objects)

This script:
  1. Walks the objects root directory.
  2. Moves any file not at the top level into <root>/<filename>.
  3. Updates storage_uri in musehub_objects for every moved file.
  4. Prints a summary.

Run inside the container or on the EC2 instance where /data is mounted:
    python3 deploy/flatten_object_store.py [--root /data/musehub/objects] [--db-url ...]

Dry-run mode (no writes):
    python3 deploy/flatten_object_store.py --dry-run
"""
from __future__ import annotations

import argparse
import os
import shutil
import sys
from pathlib import Path


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--root", default=os.environ.get("MUSEHUB_OBJECTS_DIR", "/data/musehub/objects"),
                   help="Objects root directory (musehub_objects_dir)")
    p.add_argument("--db-url", default=os.environ.get("DATABASE_URL"),
                   help="PostgreSQL connection URL for updating storage_uri")
    p.add_argument("--dry-run", action="store_true", help="Print what would happen, make no changes")
    return p.parse_args()


def flatten(root: Path, dry_run: bool) -> dict[str, str]:
    """Move all non-top-level object files to root.

    Returns {old_uri: new_uri} for every file that was moved.
    """
    moved: dict[str, str] = {}
    # Walk every file under root
    for dirpath, dirnames, filenames in os.walk(root):
        dp = Path(dirpath)
        if dp == root:
            # Top-level files are already correct — skip
            continue
        for fname in filenames:
            src = dp / fname
            dst = root / fname
            old_uri = "local://" + str(src)
            new_uri = "local://" + str(dst)
            if dst.exists():
                # Already migrated or duplicate — verify content matches then remove src
                if src.read_bytes() == dst.read_bytes():
                    print(f"  dup  {src.relative_to(root)} (already at top level)")
                    if not dry_run:
                        src.unlink()
                else:
                    print(f"  CONFLICT {src.relative_to(root)} — content differs, skipping", file=sys.stderr)
                continue
            print(f"  move {src.relative_to(root)}  →  {fname}")
            if not dry_run:
                shutil.move(str(src), str(dst))
            moved[old_uri] = new_uri

    # Remove now-empty subdirectories
    if not dry_run:
        for dirpath, dirnames, filenames in os.walk(root, topdown=False):
            dp = Path(dirpath)
            if dp == root:
                continue
            try:
                dp.rmdir()  # only removes empty dirs
                print(f"  rmdir {dp.relative_to(root)}")
            except OSError:
                pass  # not empty — leave it

    return moved


def update_db(moved: dict[str, str], db_url: str, dry_run: bool) -> None:
    """Bulk-update storage_uri in musehub_objects for all moved files."""
    if not moved:
        print("No DB updates needed.")
        return

    try:
        import psycopg2  # type: ignore[import]
    except ImportError:
        print("psycopg2 not available — skipping DB update. Run manually:", file=sys.stderr)
        for old, new in moved.items():
            print(f"  UPDATE musehub_objects SET storage_uri = '{new}' WHERE storage_uri = '{old}';")
        return

    conn = psycopg2.connect(db_url.replace("+asyncpg", "").replace("+psycopg2", ""))
    cur = conn.cursor()
    updated = 0
    for old_uri, new_uri in moved.items():
        if dry_run:
            print(f"  [dry] UPDATE storage_uri: {old_uri[:60]}... → {new_uri[:60]}...")
        else:
            cur.execute(
                "UPDATE musehub_objects SET storage_uri = %s WHERE storage_uri = %s",
                (new_uri, old_uri),
            )
            updated += cur.rowcount
    if not dry_run:
        conn.commit()
        print(f"  Updated {updated} DB row(s).")
    conn.close()


def main() -> None:
    args = parse_args()
    root = Path(args.root)

    if not root.exists():
        print(f"Root does not exist: {root}", file=sys.stderr)
        sys.exit(1)

    print(f"{'[DRY RUN] ' if args.dry_run else ''}Flattening object store: {root}")
    moved = flatten(root, args.dry_run)
    print(f"\n{'Would move' if args.dry_run else 'Moved'} {len(moved)} file(s).")

    if args.db_url:
        print("\nUpdating DB storage_uri...")
        update_db(moved, args.db_url, args.dry_run)
    else:
        print("\nNo --db-url provided — skipping DB update.")
        if moved:
            print("Run with --db-url to update storage_uri in musehub_objects.")

    print("\nDone.")


if __name__ == "__main__":
    main()