#!/usr/bin/env python3
"""Remove a background, then upscale the actual cutout through the batch API.

Python 3.10+, standard library only. Running without --run makes a local plan
and sends no requests. --run and --max-charge-usd explicitly authorize both
paid stages. A dry run precedes each submission; the cap applies to their total
quoted holds. Quotes are checked immediately before submitting, not atomically
with acceptance; the API's fare is fixed when each batch is accepted.

Progress lives in OUTPUT_DIR/workflow.json (no credentials or signed URLs).
Rerun with the same input, params and output directory to resume, including an
uncertain POST: its exact body and idempotency key are reused. A different output
directory creates a new workflow and can charge again. Keep workflow.json until
you have downloaded both results. After a timeout, work can still be running;
this program does not cancel it or start a replacement automatically.
"""

import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import struct
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener
import uuid

MAX_INPUT_BYTES = 25 * 1024 * 1024
MAX_CHAIN_PIXELS = 2_000_000
MAX_OUTPUT_BYTES = 160 * 1024 * 1024
API_TIMEOUT = 60


class WorkflowError(Exception):
    pass


class NoRedirects(HTTPRedirectHandler):
    # In particular, an API redirect must never forward its Authorization header.
    def redirect_request(self, request, fp, code, message, headers, new_url):
        return None


def image_info(data):
    """Read dimensions without decoding pixels; workers still validate the file."""
    if data.startswith(b"\x89PNG\r\n\x1a\n") and len(data) >= 24 and data[12:16] == b"IHDR":
        width, height = struct.unpack(">II", data[16:24])
        kind = "image/png"
    elif data[:2] == b"\xff\xd8":
        cursor = 2
        while cursor + 4 <= len(data):
            if data[cursor] != 255:
                raise WorkflowError("The JPEG header could not be read.")
            while cursor < len(data) and data[cursor] == 255:
                cursor += 1
            if cursor >= len(data):
                break
            marker = data[cursor]
            cursor += 1
            if marker in (0xD8, 0xD9) or 0xD0 <= marker <= 0xD7:
                continue
            if cursor + 2 > len(data):
                break
            length = int.from_bytes(data[cursor:cursor + 2], "big")
            if length < 2 or cursor + length > len(data):
                break
            if marker in (0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF) and length >= 7:
                height, width = struct.unpack(">HH", data[cursor + 3:cursor + 7])
                return checked_info(width, height, "image/jpeg")
            if marker == 0xDA:
                break
            cursor += length
        raise WorkflowError("The JPEG dimensions could not be read.")
    elif len(data) >= 25 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        chunk = data[12:16]
        if chunk == b"VP8X" and len(data) >= 30:
            if data[20] & 0x02:
                raise WorkflowError("Use a static product photo rather than an animated WebP.")
            width = 1 + int.from_bytes(data[24:27], "little")
            height = 1 + int.from_bytes(data[27:30], "little")
        elif chunk == b"VP8 " and len(data) >= 30 and data[23:26] == b"\x9d\x01\x2a":
            width, height = struct.unpack("<HH", data[26:30])
            width, height = width & 0x3FFF, height & 0x3FFF
        elif chunk == b"VP8L" and data[20] == 0x2F:
            bits = int.from_bytes(data[21:25], "little")
            width, height = 1 + (bits & 0x3FFF), 1 + ((bits >> 14) & 0x3FFF)
        else:
            raise WorkflowError("This WebP header is not supported by this example.")
        kind = "image/webp"
    else:
        raise WorkflowError("Use a PNG, JPEG, or WebP image.")
    return checked_info(width, height, kind)


def checked_info(width, height, kind):
    if width < 1 or height < 1:
        raise WorkflowError("Image dimensions must be positive.")
    return width, height, kind


def input_info(data):
    if not data or len(data) > MAX_INPUT_BYTES:
        raise WorkflowError("Each input must be nonempty and no larger than 25 MiB.")
    width, height, kind = image_info(data)
    # BiRefNet accepts more, but its output must fit Real-ESRGAN's input limit.
    if width * height > MAX_CHAIN_PIXELS:
        raise WorkflowError("This two-stage example needs an input at 2 MP or under.")
    return width, height, kind, max(1, width * height / 1_000_000)


def checked_url(url):
    parsed = urlsplit(url)
    local = parsed.hostname in ("127.0.0.1", "localhost", "::1")
    if parsed.username or parsed.password or not parsed.hostname or (parsed.scheme != "https" and not (parsed.scheme == "http" and local)):
        raise WorkflowError("Use HTTPS, or a loopback HTTP server for local testing.")
    return url


class Client:
    def __init__(self, api, key):
        self.api = checked_url(api).rstrip("/")
        self.key = key
        self.opener = build_opener(NoRedirects())

    def request(self, method, url, body=None, headers=None, limit=MAX_OUTPUT_BYTES):
        checked_url(url)
        request = Request(url, data=body, headers=headers or {}, method=method)
        try:
            with self.opener.open(request, timeout=API_TIMEOUT) as response:
                result = response.read(limit + 1)
                if len(result) > limit:
                    raise WorkflowError("The response exceeded this example's file-size limit.")
                return result
        except HTTPError as error:
            # URLs can contain signatures, and exception text can contain request data.
            code = ""
            try:
                value = json.loads(error.read(8192)).get("error", {}).get("code", "")
                visible_codes = {"unauthorized", "payment_required", "insufficient_funds", "model_unavailable", "intake_paused", "invalid_input", "rate_limited", "idempotency_conflict", "upstream_unavailable", "storage_limit", "account_suspended"}
                code = f" ({value})" if isinstance(value, str) and value in visible_codes else ""
            except (ValueError, AttributeError):
                pass
            raise WorkflowError(f"Request refused: HTTP {error.code}{code}. Keep the state file and review your account or retry later.") from None
        except (URLError, TimeoutError, OSError):
            raise WorkflowError("A network request failed. Work may still be running. Resume with the same output directory.") from None

    def api_json(self, method, path, body=None, idempotency=None):
        headers = {"Authorization": f"Bearer {self.key}"}
        if body is not None:
            headers["Content-Type"] = "application/json"
        if idempotency:
            headers["Idempotency-Key"] = idempotency
        raw = self.request(method, f"{self.api}{path}", None if body is None else json.dumps(body).encode(), headers, 4 * 1024 * 1024)
        try:
            return json.loads(raw)
        except ValueError:
            raise WorkflowError("The API returned unreadable JSON. Keep the state file and retry later.") from None

    def put(self, url, data, content_type):
        # Signed storage requests deliberately do not use api_json or its bearer header.
        self.request("PUT", url, data, {"Content-Type": content_type}, 1024 * 1024)

    def download(self, url):
        return self.request("GET", url)


def save_state(path, state):
    temp = path.with_suffix(".tmp")
    with temp.open("w", encoding="utf-8") as file:
        json.dump(state, file, indent=2)
        file.write("\n")
    temp.replace(path)


def load_state(path, expected):
    if path.exists():
        try:
            state = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            raise WorkflowError("The workflow state cannot be read. Preserve it before trying a new output directory.") from None
        if state.get("version") != 1 or any(state.get(key) != value for key, value in expected.items()):
            raise WorkflowError("This output directory belongs to another input, API, or parameter set. Use its original command to resume.")
        return state
    return {"version": 1, **expected, "workflow_id": str(uuid.uuid4()), "stages": {}}


def upload(client, data, content_type, stage, path, state):
    if "file_id" not in stage:
        result = client.api_json("POST", "/v1/files", {"files": [{"content_type": content_type, "bytes": len(data)}]})
        file = result["files"][0]
        stage["file_id"] = file["id"]
        save_state(path, state)
        url = file["upload_url"]
    elif not stage.get("uploaded"):
        result = client.api_json("POST", f"/v1/files/{stage['file_id']}/upload_url")
        url = result["upload_url"]
    else:
        return
    client.put(url, data, content_type)
    stage["uploaded"] = True
    save_state(path, state)


def run_stage(client, name, model, params, data, filename, state, state_path, cap, timeout, poll_interval):
    width, height, kind, units = input_info(data)
    digest = hashlib.sha256(data).hexdigest()
    stage = state["stages"].setdefault(name, {"key": f"{state['workflow_id']}-{name}", "input_sha256": digest})
    if stage.get("input_sha256") != digest:
        raise WorkflowError("The input for this saved stage changed. Preserve the state and original files.")
    if stage.get("output_sha256") and filename.exists():
        saved = filename.read_bytes()
        if hashlib.sha256(saved).hexdigest() != stage["output_sha256"]:
            raise WorkflowError("A saved output changed. Preserve workflow.json and restore the original output before resuming.")
        return saved
    save_state(state_path, state)
    upload(client, data, kind, stage, state_path, state)
    if "body" not in stage:
        stage["body"] = {"model": model, "items": [{"file": stage["file_id"], "params": params, "max_units": units}]}
        save_state(state_path, state)
    if "batch_id" not in stage:
        preview = client.api_json("POST", "/v1/batches", {**stage["body"], "dry_run": True})
        if preview.get("dry_run") is not True or preview.get("missing") != [] or type(preview.get("held_micros")) is not int or preview["held_micros"] < 0:
            raise WorkflowError("The dry run did not confirm a usable uploaded file and hold.")
        total = preview["held_micros"] + sum(other.get("held_micros", 0) for key, other in state["stages"].items() if key != name)
        if total > cap:
            raise WorkflowError("The two-stage quoted holds exceed --max-charge-usd. No new batch was submitted at this step.")
        stage["held_micros"] = preview["held_micros"]
        save_state(state_path, state)
        print(f"{name}: dry-run hold ${preview['held_micros'] / 1_000_000:.6f}; submitting with the saved idempotency key.")
        result = client.api_json("POST", "/v1/batches", stage["body"], stage["key"])
        if not isinstance(result.get("id"), str):
            raise WorkflowError("The API did not return a batch ID. Preserve the state and retry with this directory.")
        stage["batch_id"] = result["id"]
        save_state(state_path, state)
        accepted_hold = result.get("held_micros")
        if type(accepted_hold) is not int or accepted_hold < 0:
            raise WorkflowError("A batch was accepted without a readable hold. Its ID was saved. Check your account before resuming.")
        stage["held_micros"] = accepted_hold
        save_state(state_path, state)
        if sum(other.get("held_micros", 0) for other in state["stages"].values()) > cap:
            raise WorkflowError("A fare changed between quote and acceptance, exceeding --max-charge-usd. The accepted batch may still run; its ID was saved. No following stage was started.")
    elif sum(other.get("held_micros", 0) for other in state["stages"].values()) > cap:
        raise WorkflowError("Saved accepted holds exceed --max-charge-usd. Check your account before resuming with a different cap.")
    deadline = time.monotonic() + timeout
    while True:
        batch = client.api_json("GET", f"/v1/batches/{stage['batch_id']}")
        # Cancelled batches can still have running work: state is not the end signal.
        if batch.get("settled_at") is not None:
            break
        if time.monotonic() >= deadline:
            raise WorkflowError("Polling timed out. The batch may still be running. Resume with the same output directory.")
        time.sleep(poll_interval)
    stage["charged_micros"] = batch.get("charged_micros")
    save_state(state_path, state)
    after = None
    seen = set()
    output = None
    while True:
        params_query = {"limit": 100, "state": "succeeded"}
        if after is not None:
            params_query["after"] = after
        page = client.api_json("GET", f"/v1/batches/{stage['batch_id']}/items?{urlencode(params_query)}")
        for item in page.get("items", []):
            if item.get("state") == "succeeded" and not item.get("output_expired") and item.get("output_url"):
                output = item["output_url"]
                break
        if output or page.get("next_after") is None:
            break
        after = page["next_after"]
        if after in seen:
            raise WorkflowError("The items endpoint repeated a page cursor.")
        seen.add(after)
    if output is None:
        raise WorkflowError(f"The {name} batch has no available successful output. Check it in your account; no following stage was started.")
    result = client.download(output)
    image_info(result)
    filename.write_bytes(result)
    stage["output_sha256"] = hashlib.sha256(result).hexdigest()
    stage["output_file"] = filename.name
    save_state(state_path, state)
    print(f"{name}: downloaded {filename.name}; charged ${int(stage['charged_micros'] or 0) / 1_000_000:.6f}.")
    return result


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("input", type=Path)
    parser.add_argument("--run", action="store_true", help="Authorize both paid batches; otherwise make a local plan only")
    parser.add_argument("--scale", type=int, choices=(2, 4), default=2)
    parser.add_argument("--format", choices=("png", "webp"), default="png", help="Both formats can retain transparency")
    parser.add_argument("--output-dir", type=Path, default=Path("cutout-upscale-output"))
    parser.add_argument("--api-url", default="https://www.y-pad.com")
    parser.add_argument("--max-charge-usd", type=str, help="Required with --run; cap the sum of both quoted holds")
    parser.add_argument("--timeout", type=float, default=1800, help="Polling timeout in seconds for each stage")
    parser.add_argument("--poll-interval", type=float, default=10, help="Seconds between polling requests")
    args = parser.parse_args(argv)
    try:
        data = args.input.read_bytes()
        width, height, kind, units = input_info(data)
        checked_url(args.api_url)
        if args.timeout <= 0 or not math.isfinite(args.timeout) or args.poll_interval < 0 or not math.isfinite(args.poll_interval):
            raise WorkflowError("Timeout must be positive and polling interval nonnegative.")
        print(f"Input: {width} × {height}, {kind}; {units:g} maximum billed units per stage.")
        print(f"Plan: BiRefNet cutout → upload the cutout → Real-ESRGAN {args.scale}× → {width * args.scale} × {height * args.scale} {args.format.upper()}.")
        print("These are two submissions, each with its own published submission minimum. Review the current fare table.")
        if not args.run:
            print("Local plan only: no network requests. Add --run --max-charge-usd AMOUNT to authorize both paid stages.")
            return 0
        from decimal import Decimal, InvalidOperation
        try:
            amount = Decimal(args.max_charge_usd or "")
            if not amount.is_finite() or amount <= 0:
                raise InvalidOperation
            cap = int(amount * 1_000_000)
            if cap < 1:
                raise InvalidOperation
        except (InvalidOperation, ValueError, OverflowError):
            raise WorkflowError("--run requires a positive --max-charge-usd amount.") from None
        key = os.environ.get("YPAD_KEY", "").strip()
        if not key:
            raise WorkflowError("Set YPAD_KEY in your environment before using --run.")
        args.output_dir.mkdir(parents=True, exist_ok=True)
        state_path = args.output_dir / "workflow.json"
        state = load_state(state_path, {"input_sha256": hashlib.sha256(data).hexdigest(), "api": args.api_url.rstrip("/"), "scale": args.scale, "format": args.format})
        save_state(state_path, state)
        client = Client(args.api_url, key)
        cutout = run_stage(client, "cutout", "birefnet", {"output": "cutout"}, data, args.output_dir / "cutout.png", state, state_path, cap, args.timeout, args.poll_interval)
        run_stage(client, "upscale", "real-esrgan", {"scale": args.scale, "output_format": args.format}, cutout, args.output_dir / f"upscaled.{args.format}", state, state_path, cap, args.timeout, args.poll_interval)
        total = sum(stage.get("charged_micros") or 0 for stage in state["stages"].values())
        print(f"Complete. Total settled charges: ${total / 1_000_000:.6f}. Keep the files and state to avoid starting a replacement workflow.")
        return 0
    except KeyboardInterrupt:
        print("Stopped: work may still be running. Resume with the same input, parameters, and output directory.", file=sys.stderr)
        return 130
    except (WorkflowError, OSError) as error:
        # Only our own fixed errors are printed; filesystem errors can include private paths.
        message = str(error) if isinstance(error, WorkflowError) else "A local file could not be read or written. Check your paths and permissions."
        print(f"Stopped: {message}", file=sys.stderr)
        return 1
    except (KeyError, TypeError, ValueError, AttributeError):
        print("Stopped: an API response or saved state had an unexpected shape. Preserve workflow.json before retrying.", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())