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

import argparse
import contextlib
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import boto3

from server import (
    SETTINGS,
    analyze_uploaded_html,
    artifact_state_path,
    assemble,
    build_passthrough_upload_plan,
    build_template_plan,
    extract_body_residue_without_delivery_shells,
    load_artifact_state,
    persist_artifact_state,
    render_template_instance,
    upload_to_r2,
    validate_delivery_ready_artifact,
)


@dataclass
class RepairPreview:
    html: str
    plan: dict[str, Any]
    title: str
    warnings: list[dict[str, Any]]
    quality_score: int | None
    render_validation_summary: dict[str, Any]
    local_path: Path


def r2_client():
    if not SETTINGS.has_r2:
        return None
    return 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",
    )


def iter_saved_state_ids() -> list[str]:
    return sorted(path.parent.name for path in SETTINGS.output_dir.glob(f"*/{artifact_state_path('x').name}") if path.is_file())


def preferred_local_html_path(state: dict[str, Any]) -> Path:
    artifact_id = str(state.get("artifact_id") or "")
    current = state.get("current") if isinstance(state.get("current"), dict) else {}
    current_local_path = current.get("local_path")
    if current_local_path:
        return Path(str(current_local_path))
    html_file = str(current.get("html_file") or "").strip()
    if not html_file:
        object_key = str(current.get("object_key") or "")
        html_file = Path(object_key).name if object_key else f"{artifact_id}.html"
    return SETTINGS.output_dir / artifact_id / html_file


def fetch_current_published_html(state: dict[str, Any], client: Any | None) -> str:
    current = state.get("current") if isinstance(state.get("current"), dict) else {}
    object_key = str(current.get("object_key") or "").strip()
    if client is not None and object_key:
        body = client.get_object(Bucket=SETTINGS.r2_bucket, Key=object_key)["Body"].read()
        return body.decode("utf-8")

    local_path = preferred_local_html_path(state)
    return local_path.read_text(encoding="utf-8")


def load_admin_source_backup(state: dict[str, Any], client: Any | None) -> str | None:
    current = state.get("current") if isinstance(state.get("current"), dict) else {}
    regeneration_context = (
        current.get("regeneration_context") if isinstance(current.get("regeneration_context"), dict) else {}
    )

    local_backup = str(regeneration_context.get("admin_source_backup_local_path") or "").strip()
    if local_backup:
        local_path = Path(local_backup)
        if local_path.exists():
            return local_path.read_text(encoding="utf-8")

    artifact_id = str(state.get("artifact_id") or "")
    object_key = str(regeneration_context.get("admin_source_backup_object_key") or "").strip()
    if not object_key and artifact_id:
        object_key = f"{SETTINGS.r2_prefix}/{artifact_id}/_source/original-upload.html"

    if client is not None and object_key:
        try:
            body = client.get_object(Bucket=SETTINGS.r2_bucket, Key=object_key)["Body"].read()
            return body.decode("utf-8")
        except Exception:
            return None

    return None


def build_repair_preview(state: dict[str, Any]) -> RepairPreview:
    artifact_id = str(state.get("artifact_id") or "")
    title = str(state.get("title") or "IterFact Artifact")
    prompt = str(state.get("prompt") or title)
    template_hint = state.get("template_hint")
    theme_hint = state.get("theme_hint")
    interaction_mode = state.get("interaction_mode")
    generation_mode = str(state.get("generation_mode") or "planned_render")
    structured_data = state.get("structured_data") if isinstance(state.get("structured_data"), dict) else None
    engine_override = state.get("engine_override_requested")
    section_count_requested = state.get("section_count_requested")
    uploaded_content = state.get("uploaded_content")
    plan = state.get("plan") if isinstance(state.get("plan"), dict) else {}

    raw_html: str

    if generation_mode == "admin_passthrough_html":
        source_html = str(uploaded_content or "")
        if not source_html.strip():
            source_html = str(load_admin_source_backup(state, r2_client()) or "")
        if not source_html.strip():
            raise ValueError(f"{artifact_id} is admin_passthrough_html but has no uploaded_content to restore from.")
        if not plan:
            upload_analysis = analyze_uploaded_html(source_html, prompt)
            plan = build_passthrough_upload_plan(prompt, upload_analysis, title)
        raw_html = source_html
    elif generation_mode == "template_instance_render":
        template_id = str((structured_data or {}).get("template_id") or (template_hint or "")).strip()
        normalized_data = (structured_data or {}).get("normalized_data")
        if not template_id or not isinstance(normalized_data, dict):
            raise ValueError(f"{artifact_id} is template_instance_render but lacks template_id/normalized_data.")
        display_name = (structured_data or {}).get("display_name")
        network_label = (structured_data or {}).get("network_label")
        compose = (structured_data or {}).get("compose")
        plan = build_template_plan(
            template_id=template_id,
            normalized_data=normalized_data,
            display_name=display_name if isinstance(display_name, str) else None,
            network_label=network_label if isinstance(network_label, str) else None,
            compose=compose if isinstance(compose, list) else None,
        )
        raw_html = render_template_instance(
            template_id=template_id,
            normalized_data=normalized_data,
            display_name=display_name if isinstance(display_name, str) else None,
            network_label=network_label if isinstance(network_label, str) else None,
            compose=compose if isinstance(compose, list) else None,
        )
        title = str(plan.get("meta", {}).get("title") or title)
        theme_hint = str(plan.get("theme") or theme_hint or "")
    else:
        if not plan:
            raise ValueError(f"{artifact_id} does not have a saved plan to rebuild from.")
        raw_html = assemble(plan, str(SETTINGS.vault_path / "engines"))

    repaired_html, warnings, quality_score, render_validation_summary = validate_delivery_ready_artifact(
        html=raw_html,
        artifact_id=artifact_id,
        artifact_title=title,
        interaction_mode=interaction_mode if isinstance(interaction_mode, str) else None,
        plan=plan,
        requested_section_count=int(section_count_requested) if section_count_requested is not None else None,
        engine_override=str(engine_override) if engine_override is not None else None,
        prompt=prompt,
        template_hint=str(template_hint) if template_hint is not None else None,
        uploaded_content=str(uploaded_content) if uploaded_content is not None else None,
        structured_data=structured_data,
        generation_mode=generation_mode,
        payload=None,
    )

    return RepairPreview(
        html=repaired_html,
        plan=plan,
        title=title,
        warnings=warnings,
        quality_score=quality_score,
        render_validation_summary=render_validation_summary,
        local_path=preferred_local_html_path(state),
    )


def repair_reason(state: dict[str, Any], current_html: str, repaired_html: str) -> str | None:
    generation_mode = str(state.get("generation_mode") or "")
    current_residue = extract_body_residue_without_delivery_shells(current_html)
    repaired_residue = extract_body_residue_without_delivery_shells(repaired_html)

    if current_html.count("IFX_LIBRARY_SHELL_START") > 1 or current_html.count("IFX_NEXT_MOVE_SHELL_START") > 1:
        return "duplicate delivery shells detected"
    if repaired_residue and not current_residue:
        return "published artifact body is missing after delivery-shell stripping"
    if generation_mode == "admin_passthrough_html" and current_html != repaired_html:
        return "admin passthrough differs from canonical rebuilt upload"
    return None


def backup_corrupt_html(local_path: Path, current_html: str) -> Path:
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    backup_path = local_path.with_name(f"{local_path.stem}.pre-repair-{timestamp}.html")
    backup_path.parent.mkdir(parents=True, exist_ok=True)
    backup_path.write_text(current_html, encoding="utf-8")
    return backup_path


def apply_repair(state: dict[str, Any], preview: RepairPreview, current_html: str) -> dict[str, Any]:
    artifact_id = str(state.get("artifact_id") or "")
    local_path = preview.local_path
    local_path.parent.mkdir(parents=True, exist_ok=True)
    backup_path = backup_corrupt_html(local_path, current_html)
    local_path.write_text(preview.html, encoding="utf-8")

    current = state.get("current") if isinstance(state.get("current"), dict) else {}
    storage = str(current.get("storage") or ("r2" if SETTINGS.has_r2 else "local"))
    object_key = str(current.get("object_key") or "") or None
    preview_url = str(current.get("preview_url") or "") or None
    if SETTINGS.has_r2:
        object_key, preview_url = upload_to_r2(local_path, artifact_id)
        storage = "r2"

    persisted_state = persist_artifact_state(
        artifact_id=artifact_id,
        title=preview.title,
        prompt=str(state.get("prompt") or preview.title),
        template_hint=str(state.get("template_hint") or "") or None,
        theme_hint=str(state.get("theme_hint") or "") or None,
        interaction_mode=str(state.get("interaction_mode") or "") or None,
        generation_mode=str(state.get("generation_mode") or "planned_render"),
        structured_data=state.get("structured_data") if isinstance(state.get("structured_data"), dict) else None,
        engine_override=str(state.get("engine_override_requested") or "") or None,
        section_count_requested=int(state.get("section_count_requested")) if state.get("section_count_requested") is not None else None,
        plan=preview.plan,
        local_path=local_path,
        storage=storage,
        object_key=object_key,
        preview_url=preview_url,
        warnings=preview.warnings,
        quality_score=preview.quality_score,
        render_validation_summary=preview.render_validation_summary,
        regeneration_context=current.get("regeneration_context") if isinstance(current.get("regeneration_context"), dict) else None,
        draft_state=state.get("draft") if isinstance(state.get("draft"), dict) else None,
        draft_blueprint=state.get("draft_blueprint") if isinstance(state.get("draft_blueprint"), dict) else None,
        critic_skipped=bool(current.get("critic_skipped")),
        critic_timed_out=bool(current.get("critic_timed_out")),
        uploaded_content=str(state.get("uploaded_content")) if state.get("uploaded_content") is not None else None,
    )
    return {
        "artifact_id": artifact_id,
        "backup_path": str(backup_path),
        "preview_url": preview_url,
        "object_key": object_key,
        "updated_at": persisted_state.get("updated_at"),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Repair IterFact artifacts corrupted by delivery-shell backfill.")
    parser.add_argument("--artifact-id", action="append", dest="artifact_ids", help="Repair a specific artifact id. Can be repeated.")
    parser.add_argument("--apply", action="store_true", help="Write repaired HTML locally and push it back to R2.")
    parser.add_argument("--limit", type=int, default=None, help="Only inspect this many saved-state artifacts.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    requested_ids = list(dict.fromkeys(args.artifact_ids or []))
    artifact_ids = requested_ids or iter_saved_state_ids()
    if args.limit is not None and args.limit > 0:
        artifact_ids = artifact_ids[: args.limit]

    client = r2_client()
    inspected = 0
    repairable = 0
    repaired = 0
    skipped = 0
    failed = 0

    for artifact_id in artifact_ids:
        inspected += 1
        try:
            state = load_artifact_state(artifact_id)
            current_html = fetch_current_published_html(state, client)
            preview = build_repair_preview(state)
            reason = repair_reason(state, current_html, preview.html)
            if not reason:
                skipped += 1
                print(f"SKIP {artifact_id}: published artifact does not need this repair path")
                continue

            repairable += 1
            print(f"NEEDS_REPAIR {artifact_id}: {reason}")
            if not args.apply:
                continue

            result = apply_repair(state, preview, current_html)
            repaired += 1
            print(
                f"REPAIRED {artifact_id}: preview_url={result['preview_url']} backup={result['backup_path']}"
            )
        except FileNotFoundError as exc:
            failed += 1
            print(f"ERROR {artifact_id}: no saved state ({exc})")
        except Exception as exc:  # noqa: BLE001
            failed += 1
            print(f"ERROR {artifact_id}: {exc}")

    print(
        f"Summary: inspected={inspected} repairable={repairable} repaired={repaired} skipped={skipped} failed={failed}"
    )


if __name__ == "__main__":
    main()
