"""Migrate legacy flat-store objects to the per-repo layout.

Run in order:

    python3 migrate_flat_to_per_repo.py --dry-run   # count, touch nothing
    python3 migrate_flat_to_per_repo.py --migrate   # copy legacy → per-repo (no deletes)
    python3 migrate_flat_to_per_repo.py --verify    # confirm every ref exists in per-repo
    python3 migrate_flat_to_per_repo.py --prune     # delete legacy files (only after verify passes)

Each object is copied into EVERY repo that references it (per the musehub_object_refs table),
so the per-repo store is fully self-contained for each repo.

Source resolution per object (in order):
    1. disk_path if non-empty and is a regular file
    2. storage_uri stripped of "local://" prefix

Legacy source layouts covered:
    /data/musehub/objects/objects/<64hex>       flat store
    /data/musehub/objects/<uuid>/<64hex>        old UUID-per-repo store
    /data/musehub/sha256_<hex>                  oldest single-file layout

Target layout:
    /data/repos/<owner>/<slug>/objects/sha256/<2hex>/<62hex>

Every step is idempotent — safe to re-run. --migrate never deletes.
--prune refuses to run if --verify has not passed cleanly first.
"""
from __future__ import annotations

import argparse
import asyncio
import os
import sys
from pathlib import Path

sys.path.insert(0, "/app")
sys.path.insert(0, "/tmp/devpkgs")

REPOS_DIR = Path("/data/repos")


async def load_rows(engine) -> list:
    """Return ALL (object_id, disk_path, storage_uri, owner, slug) rows.

    One row per (object × repo) combination — objects referenced by N repos
    produce N rows so the object is copied into every referencing repo's store.
    Source fields (disk_path, storage_uri) are identical across all rows for
    the same object_id.
    """
    import sqlalchemy as sa
    from sqlalchemy.ext.asyncio import async_sessionmaker

    sf = async_sessionmaker(engine, expire_on_commit=False)
    async with sf() as s:
        result = await s.execute(sa.text(
            """
            SELECT
                obj.object_id,
                obj.disk_path,
                obj.storage_uri,
                r.owner,
                r.slug
            FROM musehub_objects obj
            JOIN musehub_object_refs ref ON ref.object_id = obj.object_id
            JOIN musehub_repos r ON r.repo_id = ref.repo_id
            ORDER BY obj.object_id, r.owner, r.slug
            """
        ))
        return result.fetchall()


def resolve_src(disk_path: str, storage_uri: str | None) -> Path | None:
    """Return the source Path for an object, or None if unresolvable."""
    if disk_path:
        p = Path(disk_path)
        if p.is_file():
            return p
    if storage_uri and storage_uri.startswith("local://"):
        p = Path(storage_uri[len("local://"):])
        if p.is_file():
            return p
    return None


def target_path(object_id: str, owner: str, slug: str) -> Path:
    from musehub.storage.backends import repo_root_for
    from muse.core.paths import server_objects_dir

    hex_str = object_id[7:]  # strip "sha256:"
    repo_root = repo_root_for(owner, slug, repos_dir=str(REPOS_DIR))
    objects_dir = server_objects_dir(repo_root)
    return objects_dir / "sha256" / hex_str[:2] / hex_str[2:]


async def cmd_dry_run(rows: list) -> None:
    would_copy = 0
    already_present = 0
    source_missing = 0

    for object_id, disk_path, storage_uri, owner, slug in rows:
        if not object_id.startswith("sha256:"):
            source_missing += 1
            continue
        src = resolve_src(disk_path, storage_uri)
        if src is None:
            print(f"  NO SOURCE: {object_id[:28]} ({owner}/{slug})")
            source_missing += 1
            continue
        tgt = target_path(object_id, owner, slug)
        if tgt.exists():
            already_present += 1
        else:
            would_copy += 1

    print("[DRY RUN]")
    print(f"  would copy:      {would_copy}")
    print(f"  already present: {already_present}")
    print(f"  source missing:  {source_missing}")


async def cmd_migrate(rows: list) -> int:
    copied = 0
    already_present = 0
    source_missing = 0
    errors = 0

    for object_id, disk_path, storage_uri, owner, slug in rows:
        if not object_id.startswith("sha256:"):
            source_missing += 1
            continue
        src = resolve_src(disk_path, storage_uri)
        if src is None:
            print(f"  NO SOURCE: {object_id[:28]} ({owner}/{slug})")
            source_missing += 1
            continue
        tgt = target_path(object_id, owner, slug)
        if tgt.exists():
            already_present += 1
            continue
        try:
            tgt.parent.mkdir(parents=True, exist_ok=True)
            tmp = tgt.with_suffix(".tmp")
            tmp.write_bytes(src.read_bytes())
            tmp.rename(tgt)
            copied += 1
        except Exception as exc:
            print(f"  ERROR {object_id[:28]} ({owner}/{slug}): {exc}")
            errors += 1

    print("\n[MIGRATE]")
    print(f"  copied:          {copied}")
    print(f"  already present: {already_present}")
    print(f"  source missing:  {source_missing}")
    print(f"  errors:          {errors}")
    return errors


async def cmd_verify(rows: list) -> int:
    ok = 0
    missing = 0

    seen_missing: set[tuple[str, str, str]] = set()
    for object_id, disk_path, storage_uri, owner, slug in rows:
        if not object_id.startswith("sha256:"):
            continue
        tgt = target_path(object_id, owner, slug)
        if tgt.exists():
            ok += 1
        else:
            key = (object_id, owner, slug)
            if key not in seen_missing:
                print(f"  MISSING: {object_id[:28]} ({owner}/{slug})")
                seen_missing.add(key)
            missing += 1

    print("\n[VERIFY]")
    print(f"  ok:      {ok}")
    print(f"  missing: {missing}")
    if missing == 0:
        print("  ✅ All objects present in per-repo store.")
    else:
        print(f"  ❌ {missing} object(s) not in per-repo store — run --migrate first.")
    return missing


async def cmd_prune(rows: list) -> None:
    # Collect unique source paths across all rows for the same object
    seen_src: set[Path] = set()
    for object_id, disk_path, storage_uri, owner, slug in rows:
        if not object_id.startswith("sha256:"):
            continue
        src = resolve_src(disk_path, storage_uri)
        if src is not None:
            seen_src.add(src)

    deleted = 0
    already_gone = 0
    errors = 0

    for src in seen_src:
        if not src.exists():
            already_gone += 1
            continue
        try:
            src.unlink()
            deleted += 1
        except Exception as exc:
            print(f"  ERROR deleting {src}: {exc}")
            errors += 1

    print("\n[PRUNE]")
    print(f"  deleted:      {deleted}")
    print(f"  already gone: {already_gone}")
    print(f"  errors:       {errors}")
    if errors == 0:
        print("\n  Legacy source files removed.")
        print("  You can now: rm -rf /data/musehub/objects")


async def main(args: argparse.Namespace) -> None:
    import sqlalchemy as sa
    from sqlalchemy.ext.asyncio import create_async_engine

    db_url = os.environ["DATABASE_URL"]
    engine = create_async_engine(db_url, echo=False)

    rows = await load_rows(engine)
    print(f"Object×repo ref pairs: {len(rows)}\n")

    if args.dry_run:
        await cmd_dry_run(rows)
    elif args.migrate:
        errors = await cmd_migrate(rows)
        if errors:
            sys.exit(1)
    elif args.verify:
        missing = await cmd_verify(rows)
        if missing:
            sys.exit(1)
    elif args.prune:
        missing = await cmd_verify(rows)
        if missing:
            print("\n❌ Aborting prune — verify failed. Run --migrate first.")
            sys.exit(1)
        await cmd_prune(rows)

    await engine.dispose()


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--dry-run", action="store_true", help="Count, touch nothing")
    g.add_argument("--migrate", action="store_true", help="Copy legacy → per-repo (no deletes)")
    g.add_argument("--verify", action="store_true", help="Confirm all refs exist in per-repo")
    g.add_argument("--prune", action="store_true",
                   help="Delete legacy files (runs verify first, aborts if not clean)")
    return p.parse_args()


if __name__ == "__main__":
    asyncio.run(main(parse_args()))