"""Run a recorded sample through the public inference pipeline."""

import argparse
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--text-encoder", required=True)
    parser.add_argument(
        "--case",
        choices=("product", "portrait", "holdout-detail", "edit"),
        default="product",
    )
    parser.add_argument(
        "--image", help="Original edit reference, required for the edit case"
    )
    parser.add_argument("--cache-dir", default="cache")
    parser.add_argument("--output", default="outputs/sample.png")
    args = parser.parse_args()
    if args.case == "edit" and args.image is None:
        parser.error(
            "The edit sample requires --image pointing to the original reference"
        )
    catalog = json.loads(Path(__file__).with_name("catalog.json").read_text())
    case = next(e for e in catalog["examples"] if e["case"] == args.case)

    from h3_mlx.runtime import Pipeline

    pipe = Pipeline(
        args.model,
        args.text_encoder,
        cache_dir=args.cache_dir,
        compile_blocks=True,
        eval_every=0,
        prefetch_blocks=4,
        vae_compile=True,
        vae_batch=4,
        kernel="metal-qk-exact",
        tune_matmul=True,
    )
    try:
        _, record = pipe.generate(
            case["prompt"],
            seed=case["seed"],
            width=case["requested_size"][0],
            height=case["requested_size"][1],
            images=[args.image] if args.case == "edit" else [],
            output=args.output,
            save_trace=True,
        )
        print(
            json.dumps(
                {"output": record["output"], "png_sha256": record["png_sha256"]},
                indent=2,
            )
        )
    finally:
        pipe.clear()


if __name__ == "__main__":
    main()