#!/usr/bin/env python3
from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

from server import (
    SETTINGS,
    guess_upload_content_type,
    make_r2_client,
)


@dataclass
class BackupStats:
    scanned: int = 0
    created: int = 0
    existing: int = 0
    failed: int = 0


def backup_key_for(object_key: str) -> str:
    parts = object_key.split("/")
    if len(parts) < 3:
        return ""
    artifact_id = parts[-2]
    file_name = parts[-1]
    return f"{SETTINGS.r2_prefix}/{artifact_id}/_source/{file_name}"


def object_exists(client, key: str) -> bool:
    try:
        client.head_object(Bucket=SETTINGS.r2_bucket, Key=key)
        return True
    except Exception:
        return False


def backfill_r2_html_backups() -> BackupStats:
    stats = BackupStats()
    if not SETTINGS.has_r2:
        return stats

    client = make_r2_client()
    prefix = f"{SETTINGS.r2_prefix}/"
    continuation_token: str | None = None

    while True:
        kwargs: dict[str, object] = {
            "Bucket": SETTINGS.r2_bucket,
            "Prefix": prefix,
            "MaxKeys": 1000,
        }
        if continuation_token:
            kwargs["ContinuationToken"] = continuation_token

        response = client.list_objects_v2(**kwargs)
        for item in response.get("Contents", []):
            key = str(item["Key"])
            if not key.endswith(".html") or "/_source/" in key or "/_versions/" in key:
                continue

            stats.scanned += 1
            backup_key = backup_key_for(key)
            if not backup_key:
                stats.failed += 1
                continue
            if object_exists(client, backup_key):
                stats.existing += 1
                continue

            try:
                body = client.get_object(Bucket=SETTINGS.r2_bucket, Key=key)["Body"].read()
                client.put_object(
                    Bucket=SETTINGS.r2_bucket,
                    Key=backup_key,
                    Body=body,
                    ContentType=guess_upload_content_type(Path(key)),
                )
                stats.created += 1
            except Exception:
                stats.failed += 1

        if not response.get("IsTruncated"):
            break
        continuation_token = response.get("NextContinuationToken")

    return stats


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill immutable _source HTML backups for hosted IterFact artifacts.",
    )
    parser.parse_args()
    stats = backfill_r2_html_backups()
    print(
        f"R2 HTML backups: scanned={stats.scanned} created={stats.created} "
        f"existing={stats.existing} failed={stats.failed}"
    )


if __name__ == "__main__":
    main()
