#!/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 base64
import mimetypes
import os
import sys

import requests

API_URL = "https://openrouter.ai/api/v1/images"
DEFAULT_MODEL = os.environ.get("MODEL", "meta/muse-image")

BACKGROUND_PRESETS = {
    "beige": "Soft warm beige studio background, evenly lit, no shadows, no props.",
    "white": "Plain bright white studio backdrop, evenly lit, no shadows, no gradients, no props.",
}

EXTENSIONS = {
    "image/png": ".png",
    "image/jpeg": ".jpg",
    "image/webp": ".webp",
}


def build_prompt(bg_choice: str) -> str:
    bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice)
    return (
        "Product photography edit of this garment photo. Keep the garments pixel-faithful: identical "
        "shape, colors, prints, labels, buttons, and every existing flaw unchanged. Remove ALL wrinkles "
        "and creases completely: the fabric must look perfectly smooth and freshly ironed, flat like a "
        f"new catalog product photo. {bg_description}"
    )


def encode_image(path: str) -> str:
    mime = mimetypes.guess_type(path)[0] or "image/png"
    with open(path, "rb") as f:
        return f"data:{mime};base64,{base64.b64encode(f.read()).decode()}"


def restore_resolution(out: str, src: str) -> None:
    """Resize the generated image back to the input dimensions using Lanczos resampling."""
    try:
        from PIL import Image
    except ImportError:
        print("Pillow not installed; skipping resolution restore (pip install Pillow)")
        return

    with Image.open(src) as orig, Image.open(out) as gen:
        if gen.size == orig.size:
            return
        resized = gen.resize(orig.size, Image.LANCZOS)
        ext = os.path.splitext(out)[1].lower()
        fmt = "JPEG" if ext in (".jpg", ".jpeg") else None
        kwargs = {"quality": 95} if fmt == "JPEG" else {}
        if fmt:
            resized.save(out, fmt, **kwargs)
        else:
            resized.save(out)
        print(f"Restored resolution: {gen.size} -> {orig.size}")


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}...")
    resp = requests.post(
        API_URL,
        headers={"Authorization": f"Bearer {args.api_key}"},
        json={
            "model": args.model,
            "prompt": prompt,
            "background": "opaque",
            "input_references": [
                {"type": "image_url", "image_url": {"url": encode_image(args.image)}}
            ],
            "output_format": "jpeg",
        },
        timeout=300,
    )

    if resp.status_code != 200:
        try:
            message = resp.json()["error"]["message"]
        except Exception:
            message = resp.text
        sys.exit(f"API error {resp.status_code}: {message}")

    result = resp.json()
    image = result["data"][0]
    ext = EXTENSIONS.get(image.get("media_type"), ".jpg")

    out = args.output or (os.path.splitext(args.image)[0] + "_clean" + ext)

    with open(out, "wb") as f:
        f.write(base64.b64decode(image["b64_json"]))

    if not args.no_restore_res:
        restore_resolution(out, args.image)

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


if __name__ == "__main__":
    main()