#!/usr/bin/env python3
"""Unwrinkle clothing photos and unify backgrounds for marketplace listings.

Preserves garment shape, colors, prints, and defects while removing creases
and replacing the background with a studio backdrop.
"""

import argparse
import os
import sys

from cleaner import (
    API_URL,
    BACKGROUND_PRESETS,
    CleanerError,
    CleanResult,
    DEFAULT_MODEL,
    EXTENSIONS,
    build_prompt,
    clean_garment,
    encode_image,
    encode_image_bytes,
    restore_resolution,
    restore_resolution_bytes,
)

__all__ = [
    "API_URL",
    "DEFAULT_MODEL",
    "BACKGROUND_PRESETS",
    "EXTENSIONS",
    "build_prompt",
    "encode_image",
    "encode_image_bytes",
    "restore_resolution",
    "restore_resolution_bytes",
    "clean_garment",
    "CleanResult",
    "CleanerError",
    "parse_args",
    "main",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Unwrinkle clothing photos and unify backgrounds via OpenRouter."
    )
    parser.add_argument("image", help="Path to input garment photo")
    parser.add_argument(
        "-o",
        "--output",
        help="Path for cleaned output image (defaults to <name>_clean.<ext>)",
    )
    parser.add_argument(
        "--bg",
        default="beige",
        help="Background preset ('beige', 'white') or a custom background description",
    )
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"OpenRouter image model slug (default: {DEFAULT_MODEL})",
    )
    parser.add_argument(
        "--api-key",
        default=os.environ.get("OPENROUTER_API_KEY"),
        help="OpenRouter API key (defaults to OPENROUTER_API_KEY env var)",
    )
    parser.add_argument(
        "--no-restore-res",
        action="store_true",
        help="Skip upscaling back to source photo dimensions",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if not os.path.isfile(args.image):
        sys.exit(f"Error: file not found: {args.image}")

    if not args.api_key:
        sys.exit(
            "Error: missing API key. Set OPENROUTER_API_KEY in your environment or pass --api-key."
        )

    prompt = build_prompt(args.bg)

    print(f"Submitting to {args.model}...")
    try:
        result = clean_garment(
            image_data=args.image,
            prompt=prompt,
            api_key=args.api_key,
            model=args.model,
            restore_res=not args.no_restore_res,
        )
    except CleanerError as e:
        sys.exit(str(e))

    ext = EXTENSIONS.get(result.media_type, ".jpg")
    out = args.output or (os.path.splitext(args.image)[0] + "_clean" + ext)

    with open(out, "wb") as f:
        f.write(result.image_bytes)

    if not args.no_restore_res and result.was_rescaled:
        print(f"Restored resolution: {result.model_dimensions} -> {result.original_dimensions}")

    cost = result.cost
    cost_str = f" (cost ${cost:.2f})" if cost is not None else ""
    print(f"Saved: {out}{cost_str}")


if __name__ == "__main__":
    main()