#!/usr/bin/env python3
"""Server-side S3 remap: <bucket>/<name> -> stub/<bucket>/<name>/<version>.

Much faster than shelling out to the AWS CLI per object.
Uses the same TSV manifest as remap_storage_s3_keys.sh:

  ~/migration/remap_keys.tsv
  columns: src_key<TAB>dst_key

Usage on EC2:
  sudo apt-get install -y python3-boto3
  nohup env JOBS=128 python3 ~/migration/remap_storage_fast.py \\
    > ~/migration/remap-fast.out 2>&1 &
  tail -f ~/migration/remap-fast.out
"""
from __future__ import annotations

import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError

MANIFEST = os.path.expanduser(os.environ.get("MANIFEST", "~/migration/remap_keys.tsv"))
BUCKET = os.environ.get("S3_BUCKET", "datafy-prod-supabase-storage")
REGION = os.environ.get("AWS_REGION", "eu-west-1")
JOBS = int(os.environ.get("JOBS", "128"))
LOG = os.path.expanduser(os.environ.get("LOG", "~/migration/remap_fast.log"))
# How many in-flight futures to keep queued (not all 220k at once).
BATCH = int(os.environ.get("BATCH", str(JOBS * 8)))

cfg = Config(
    max_pool_connections=JOBS + 32,
    retries={"max_attempts": 10, "mode": "adaptive"},
)
s3 = boto3.client("s3", region_name=REGION, config=cfg)


def _is_missing(exc: ClientError) -> bool:
    code = exc.response.get("Error", {}).get("Code", "")
    return code in ("404", "NoSuchKey", "NotFound", "404 Not Found")


def one(src: str, dst: str) -> tuple[str, str]:
    try:
        s3.head_object(Bucket=BUCKET, Key=dst)
        return "SKIP", dst
    except ClientError as e:
        if not _is_missing(e):
            return "FAIL", f"{src} head_dst {e}"

    try:
        s3.copy_object(
            Bucket=BUCKET,
            Key=dst,
            CopySource={"Bucket": BUCKET, "Key": src},
            MetadataDirective="COPY",
        )
        return "OK", src
    except ClientError as e:
        if _is_missing(e):
            return "MISSING", src
        return "FAIL", f"{src} {e}"


def main() -> None:
    rows: list[tuple[str, str]] = []
    with open(MANIFEST, encoding="utf-8") as f:
        for line in f:
            line = line.rstrip("\n")
            if not line:
                continue
            src, dst = line.split("\t", 1)
            rows.append((src, dst))

    total = len(rows)
    ok = skip = miss = fail = 0
    print(f"start total={total} jobs={JOBS} batch={BATCH} bucket={BUCKET}", flush=True)
    t0 = time.time()
    n = 0

    with open(LOG, "a", encoding="utf-8") as log, ThreadPoolExecutor(max_workers=JOBS) as ex:
        for i in range(0, total, BATCH):
            chunk = rows[i : i + BATCH]
            futs = [ex.submit(one, s, d) for s, d in chunk]
            for fut in as_completed(futs):
                status, msg = fut.result()
                n += 1
                if status == "OK":
                    ok += 1
                elif status == "SKIP":
                    skip += 1
                elif status == "MISSING":
                    miss += 1
                else:
                    fail += 1
                log.write(f"{status} {msg}\n")
                if n % 500 == 0 or n == total:
                    elapsed = time.time() - t0
                    rate = n / elapsed if elapsed else 0.0
                    eta_h = ((total - n) / rate / 3600.0) if rate else 0.0
                    print(
                        f"progress {n}/{total} ({100.0 * n / total:.1f}%) "
                        f"ok={ok} skip={skip} miss={miss} fail={fail} "
                        f"rate={rate:.0f}/s eta={eta_h:.1f}h",
                        flush=True,
                    )

    print(f"DONE ok={ok} skip={skip} miss={miss} fail={fail}", flush=True)


if __name__ == "__main__":
    main()