import argparse
import json
import shutil
from pathlib import Path

from gradio_client import Client, handle_file


VALID_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}


def iter_images(root: Path):
    for path in sorted(root.rglob("*")):
        if path.is_file() and path.suffix.lower() in VALID_SUFFIXES:
            yield path


def main():
    parser = argparse.ArgumentParser(description="Batch client for the Catalog Creator Space.")
    parser.add_argument("--space", required=True, help="Space id, for example username/CatalogCreator")
    parser.add_argument("--input-dir", required=True, help="Directory of source images")
    parser.add_argument("--output-dir", required=True, help="Directory for processed files")
    parser.add_argument("--mode", choices=["extract_outfit", "cutout"], default="extract_outfit")
    parser.add_argument(
        "--extract-prompt",
        default="Extract the clothing and create a flat mockup.",
        help="Prompt used only for extract_outfit mode",
    )
    parser.add_argument("--output-size", type=int, default=1400)
    parser.add_argument("--padding-percent", type=int, default=8)
    parser.add_argument("--alpha-threshold", type=int, default=8)
    parser.add_argument("--num-inference-steps", type=int, default=24)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--skip-existing", action="store_true")
    args = parser.parse_args()

    input_dir = Path(args.input_dir).resolve()
    output_dir = Path(args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    client = Client(args.space)
    manifest_path = output_dir / "manifest.jsonl"

    for image_path in iter_images(input_dir):
        relative_parent = image_path.parent.relative_to(input_dir)
        target_dir = output_dir / relative_parent
        target_dir.mkdir(parents=True, exist_ok=True)

        transparent_target = target_dir / f"{image_path.stem}_transparent.png"
        square_target = target_dir / f"{image_path.stem}_catalog.png"
        preview_target = target_dir / f"{image_path.stem}_preview.png"
        metadata_target = target_dir / f"{image_path.stem}.json"

        if args.skip_existing and transparent_target.exists() and square_target.exists() and metadata_target.exists():
            continue

        try:
            _, _, transparent_file, square_file, preview_file, metadata = client.predict(
                image=handle_file(str(image_path)),
                mode=args.mode,
                output_size=args.output_size,
                padding_percent=args.padding_percent,
                alpha_threshold=args.alpha_threshold,
                crop_to_subject=True,
                extract_prompt=args.extract_prompt,
                num_inference_steps=args.num_inference_steps,
                seed=args.seed,
                api_name="/catalog_prep",
            )

            shutil.copy2(transparent_file, transparent_target)
            shutil.copy2(square_file, square_target)
            shutil.copy2(preview_file, preview_target)
            metadata_target.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

            with manifest_path.open("a", encoding="utf-8") as manifest:
                manifest.write(
                    json.dumps(
                        {
                            "source": str(image_path),
                            "transparent_png": str(transparent_target),
                            "square_png": str(square_target),
                            "preview_png": str(preview_target),
                            "metadata_json": str(metadata_target),
                            "status": "ok",
                        }
                    )
                    + "\n"
                )
        except Exception as exc:
            with manifest_path.open("a", encoding="utf-8") as manifest:
                manifest.write(
                    json.dumps(
                        {
                            "source": str(image_path),
                            "status": "error",
                            "error": str(exc),
                        }
                    )
                    + "\n"
                )


if __name__ == "__main__":
    main()