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

import argparse
import tempfile
from dataclasses import dataclass
from pathlib import Path

import boto3

from server import (
    SETTINGS,
    ensure_library_preview_assets,
    extract_body_residue_without_delivery_shells,
    guess_upload_content_type,
)


@dataclass
class PreviewBackfillStats:
    scanned: int = 0
    generated: int = 0
    fallback_generated: int = 0
    shell_only_skipped: int = 0
    failed: int = 0


def has_meaningful_preview_body(html: str) -> bool:
    return bool(extract_body_residue_without_delivery_shells(html))


def upload_sidecars(client: boto3.client, artifact_id: str, sidecars: list[Path]) -> None:
    for sidecar in sidecars:
        sidecar_key = f"{SETTINGS.r2_prefix}/{artifact_id}/{sidecar.name}"
        client.upload_file(
            str(sidecar),
            SETTINGS.r2_bucket,
            sidecar_key,
            ExtraArgs={"ContentType": guess_upload_content_type(sidecar)},
        )


def backfill_local_previews() -> PreviewBackfillStats:
    stats = PreviewBackfillStats()
    for html_path in SETTINGS.output_dir.glob("*/*.html"):
        stats.scanned += 1
        try:
            html = html_path.read_text(encoding="utf-8")
            generated = ensure_library_preview_assets(html_path, html)
            if generated:
                stats.generated += 1
                if not has_meaningful_preview_body(html):
                    stats.fallback_generated += 1
        except Exception:
            stats.failed += 1
    return stats


def backfill_r2_previews() -> PreviewBackfillStats:
    stats = PreviewBackfillStats()
    if not SETTINGS.has_r2:
        return stats

    client = boto3.client(
        "s3",
        endpoint_url=f"https://{SETTINGS.r2_account_id}.r2.cloudflarestorage.com",
        aws_access_key_id=SETTINGS.r2_access_key_id,
        aws_secret_access_key=SETTINGS.r2_secret_access_key,
        region_name="auto",
    )

    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
            parts = key.split("/")
            if len(parts) < 3:
                stats.failed += 1
                continue

            artifact_id = parts[-2]
            try:
                body = client.get_object(Bucket=SETTINGS.r2_bucket, Key=key)["Body"].read()
                html = body.decode("utf-8")

                with tempfile.TemporaryDirectory(prefix=f"ifx_preview_{artifact_id}_") as tmp:
                    local_html = Path(tmp) / Path(key).name
                    local_html.write_text(html, encoding="utf-8")
                    generated = ensure_library_preview_assets(local_html, html)
                    if not generated:
                        continue
                    upload_sidecars(client, artifact_id, generated)
                    stats.generated += 1
                    if not has_meaningful_preview_body(html):
                        stats.fallback_generated += 1
            except Exception:
                stats.failed += 1

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

    return stats


def print_stats(label: str, stats: PreviewBackfillStats) -> None:
    print(
        f"{label}: scanned={stats.scanned} generated={stats.generated} "
        f"fallback_generated={stats.fallback_generated} "
        f"shell_only_skipped={stats.shell_only_skipped} failed={stats.failed}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill cached library preview assets for existing IterFact artifacts.",
    )
    parser.add_argument(
        "--local-only",
        action="store_true",
        help="Only generate sidecars for local generated artifacts.",
    )
    args = parser.parse_args()

    local_stats = backfill_local_previews()
    print_stats("Local previews", local_stats)

    if args.local_only:
        return

    r2_stats = backfill_r2_previews()
    print_stats("R2 previews", r2_stats)


if __name__ == "__main__":
    main()
