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

import argparse
import contextlib
import re
from pathlib import Path

import boto3

from server import (
    SETTINGS,
    apply_library_shell,
    apply_next_move_shell,
    assert_delivery_shell_preserves_body,
    backup_existing_r2_object,
    load_artifact_state,
    make_r2_client,
)


TITLE_RE = re.compile(r"<title>(.*?)</title>", re.IGNORECASE | re.DOTALL)


def extract_title(html: str, fallback: str) -> str:
    match = TITLE_RE.search(html)
    if not match:
        return fallback
    title = re.sub(r"\s+", " ", match.group(1)).strip()
    return title or fallback


def patch_delivery_shells(html: str, artifact_id: str, title: str) -> str:
    state: dict[str, object] = {}
    with contextlib.suppress(Exception):
        state = load_artifact_state(artifact_id)

    generation_mode = str(state.get("generation_mode") or "")
    if generation_mode == "admin_passthrough_html":
        return html

    prompt = str(state.get("prompt") or title)
    template_hint = state.get("template_hint")
    generation_mode = generation_mode or "planned_render"
    plan = state.get("plan") if isinstance(state.get("plan"), dict) else {
        "wrapper": {},
        "meta": {},
        "sections": [],
    }

    patched = apply_library_shell(html, artifact_id, title)
    patched = apply_next_move_shell(
        patched,
        artifact_id=artifact_id,
        artifact_title=title,
        prompt=prompt,
        template_hint=template_hint if isinstance(template_hint, str) else None,
        generation_mode=generation_mode,
        plan=plan,
    )
    assert_delivery_shell_preserves_body(html, patched, artifact_id=artifact_id)
    return patched


def patch_local_artifacts() -> tuple[int, int]:
    scanned = 0
    updated = 0
    for html_path in SETTINGS.output_dir.glob("*/*.html"):
        scanned += 1
        html = html_path.read_text(encoding="utf-8")
        artifact_id = html_path.parent.name
        title = extract_title(html, "IterFact Artifact")
        patched = patch_delivery_shells(html, artifact_id, title)
        if patched == html:
            continue
        html_path.write_text(patched, encoding="utf-8")
        updated += 1
    return scanned, updated


def patch_r2_artifacts() -> tuple[int, int]:
    if not SETTINGS.has_r2:
        return 0, 0

    client = make_r2_client()

    prefix = f"{SETTINGS.r2_prefix}/"
    continuation_token: str | None = None
    scanned = 0
    updated = 0

    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 = item["Key"]
            if not str(key).endswith(".html") or "/_source/" in str(key) or "/_versions/" in str(key):
                continue

            scanned += 1
            parts = str(key).split("/")
            if len(parts) < 3:
                continue

            artifact_id = parts[-2]
            body = client.get_object(Bucket=SETTINGS.r2_bucket, Key=key)["Body"].read()
            html = body.decode("utf-8")
            title = extract_title(html, "IterFact Artifact")
            patched = patch_delivery_shells(html, artifact_id, title)
            if patched == html:
                continue

            backup_existing_r2_object(client, str(key))
            client.put_object(
                Bucket=SETTINGS.r2_bucket,
                Key=key,
                Body=patched.encode("utf-8"),
                ContentType="text/html; charset=utf-8",
            )
            updated += 1

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

    return scanned, updated


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill IterFact delivery shells into local and hosted artifacts.",
    )
    parser.add_argument(
        "--local-only",
        action="store_true",
        help="Patch only local generated artifacts and skip R2.",
    )
    args = parser.parse_args()

    local_scanned, local_updated = patch_local_artifacts()
    print(f"Local artifacts scanned: {local_scanned}; updated: {local_updated}")

    if args.local_only:
        return

    r2_scanned, r2_updated = patch_r2_artifacts()
    print(f"R2 artifacts scanned: {r2_scanned}; updated: {r2_updated}")


if __name__ == "__main__":
    main()
