#!/usr/bin/env python3
"""Run one Mage-Flow Edit XPO3 image edit."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import sys
import time

from PIL import Image


ROOT = Path(__file__).resolve().parent
for path in (ROOT / "runtime", ROOT / "vendor"):
    sys.path.insert(0, str(path))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("reference", type=Path)
    parser.add_argument("instruction")
    parser.add_argument("--output", type=Path, default=Path("edited.png"))
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--max-size", type=int, default=1024)
    parser.add_argument("--disable-fused-gelu-up", action="store_true")
    parser.add_argument("--disable-fp4-bridge", action="store_true")
    parser.add_argument("--disable-accelerated-attention", action="store_true")
    parser.add_argument("--disable-direct-hnd", action="store_true")
    parser.add_argument("--report", type=Path)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

    import torch
    from portable_turbo_runtime import (
        close_pipeline_optimization_runtimes,
        generation_optimization_context,
        load_pipeline_from_files,
        transformer_config,
    )

    if not torch.cuda.is_available():
        raise RuntimeError("XPO3 requires an NVIDIA CUDA GPU")
    torch.cuda.set_device(0)
    properties = torch.cuda.get_device_properties(0)
    if (properties.major, properties.minor) != (12, 0):
        raise RuntimeError(
            "this XPO3 build requires an NVIDIA Blackwell SM120 GPU; "
            f"found compute capability {properties.major}.{properties.minor}"
        )

    diffusion_model = (
        ROOT / "Mage-Flow-Edit-XPO3-NVFP4.safetensors"
    )
    text_encoder = ROOT / "qwen3vl_4b_fp8_scaled.safetensors"
    vae = ROOT / "Mage-Flow-VAE.safetensors"
    config = transformer_config(diffusion_model)
    profile = config["quantization_config"]["xpo3_runtime_profile"]
    bridge_blocks = ",".join(
        sorted(
            profile["fp4_bridge_scales"],
            key=lambda value: int(value),
        )
    )
    attention_steps = ",".join(
        str(value) for value in profile["attention_steps"]
    )
    attention_blocks = ",".join(
        str(value) for value in profile["attention_blocks"]
    )

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.cuda.reset_peak_memory_stats(0)
    started = time.perf_counter()
    pipe, load_report = load_pipeline_from_files(
        diffusion_model=diffusion_model,
        text_encoder=text_encoder,
        vae=vae,
        support_root=ROOT / "resources" / "edit",
        fused_gelu_library=ROOT / "runtime/libmage_nvfp4_gelu_up.so",
        bridge_up_library=ROOT / "runtime/libmage_nvfp4_bridge_up.so",
        bridge_down_library=(
            ROOT / "runtime/libmage_nvfp4_prequantized_down.so"
        ),
        torch=torch,
    )
    for module in (
        pipe.model.txt_enc,
        pipe.model.transformer,
        pipe.model.vae,
    ):
        module.to("cuda:0")
    torch.cuda.synchronize()
    load_seconds = time.perf_counter() - started

    reference = Image.open(args.reference).convert("RGB")
    generation_started = time.perf_counter()
    with torch.inference_mode():
        with generation_optimization_context(
            pipe=pipe,
            torch=torch,
            enable_fused_gelu_up=not args.disable_fused_gelu_up,
            enable_fp4_bridge=not args.disable_fp4_bridge,
            bridge_blocks=bridge_blocks,
            enable_attention_accel=(
                not args.disable_accelerated_attention
            ),
            enable_direct_hnd=not args.disable_direct_hnd,
            attention_steps=attention_steps,
            attention_blocks=attention_blocks,
            steps=int(profile["steps"]),
            static_shift=float(profile["static_shift"]),
            cfg=float(profile["cfg"]),
            required_cfg=float(profile["cfg"]),
            expected_steps=int(profile["steps"]),
        ) as feature_manifest:
            images = pipe.edit(
                [args.instruction],
                [[reference]],
                neg_prompts=[" "],
                seeds=[args.seed],
                steps=int(profile["steps"]),
                cfg=float(profile["cfg"]),
                max_size=args.max_size,
                static_shift=float(profile["static_shift"]),
                prompt_template="mage-flow-edit",
                vl_cond_long_edge=384,
                batch_cfg=False,
            )
    torch.cuda.synchronize()
    generation_seconds = time.perf_counter() - generation_started
    if not feature_manifest["restoration"]["all_restored"]:
        raise RuntimeError("XPO3 feature state did not restore after editing")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    images[0].save(args.output)
    report = {
        "schema_version": "mage-flow-edit-xpo3-cli-v1",
        "reference": str(args.reference.resolve()),
        "instruction": args.instruction,
        "output": str(args.output.resolve()),
        "seed": args.seed,
        "max_size": args.max_size,
        "gpu": {
            "name": properties.name,
            "compute_capability": (
                f"{properties.major}.{properties.minor}"
            ),
        },
        "load_and_placement_seconds": load_seconds,
        "generation_seconds": generation_seconds,
        "peak_allocated_bytes": int(torch.cuda.max_memory_allocated(0)),
        "load": load_report,
        "active_feature_manifest": feature_manifest,
    }
    if args.report is not None:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(
            json.dumps(report, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
    print(json.dumps(report, indent=2, sort_keys=True))
    cleanup = close_pipeline_optimization_runtimes(pipe)
    if not cleanup["all_closed_without_error"]:
        raise RuntimeError(f"XPO3 native cleanup failed: {cleanup}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())