#!/usr/bin/env python3
"""Python 3.10+, standard library, MIT. Offline unless --run is supplied."""
import argparse
import base64
import csv
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import re
import struct
import sys
from urllib.error import HTTPError, URLError
from urllib.request import HTTPRedirectHandler, Request, build_opener
import uuid

API = "https://api.quickrouter.ai/v1/images/edits"
MAX_INPUT = 10 * 1024 * 1024
MAX_RESPONSE = 45 * 1024 * 1024
MAX_OUTPUT = 32 * 1024 * 1024


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class RequestFailure(Exception):
    def __init__(self, code="network_or_timeout"):
        self.code = code


def digest(data):
    return hashlib.sha256(data).hexdigest()


def utc_now():
    return datetime.now(timezone.utc).isoformat()


def png_size(data):
    if len(data) < 33 or data[:8] != b"\x89PNG\r\n\x1a\n" or data[12:16] != b"IHDR":
        raise ValueError("A PNG signature and IHDR header are required")
    width, height = struct.unpack(">II", data[16:24])
    if not width or not height or width * height > 80_000_000:
        raise ValueError("PNG dimensions exceed this example's limit")
    return [width, height]


def private_write(path, data):
    if path.is_symlink():
        raise ValueError("Output symlinks are not allowed")
    temp = path.with_name(path.name + "." + uuid.uuid4().hex + ".tmp")
    try:
        fd = os.open(temp, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
        with os.fdopen(fd, "wb") as stream:
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temp, path)
    finally:
        if temp.exists():
            temp.unlink()


def save_state(path, state):
    private_write(path, json.dumps(state, ensure_ascii=False, indent=2).encode("utf-8"))


def load_state(path):
    if path.is_symlink():
        raise ValueError("Manifest symlinks are not allowed")
    if not path.exists():
        return {"version": 1, "jobs": {}}
    if path.stat().st_size > 2 * 1024 * 1024:
        raise ValueError("Manifest is too large")
    state = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(state, dict) or state.get("version") != 1 or not isinstance(state.get("jobs"), dict):
        raise ValueError("Unsupported manifest")
    for job_id, record in state["jobs"].items():
        if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]{0,63}", job_id) or not isinstance(record, dict):
            raise ValueError("Invalid manifest job")
        if not isinstance(record.get("plan_hash"), str) or not re.fullmatch(r"[a-f0-9]{64}", record["plan_hash"]):
            raise ValueError("Invalid manifest fingerprint")
        if record.get("status") not in {"indeterminate", "response_received", "succeeded"}:
            raise ValueError("Invalid manifest status")
        if not isinstance(record.get("attempts"), int) or isinstance(record["attempts"], bool) or record["attempts"] < 1:
            raise ValueError("Invalid manifest attempt count")
    return state


def load_jobs(csv_path, input_root, model, size, quality):
    root = input_root.resolve(strict=True)
    if not root.is_dir():
        raise ValueError("Input root must be a directory")
    if csv_path.stat().st_size > 2 * 1024 * 1024:
        raise ValueError("CSV is too large")
    jobs = []
    ids = set()
    with csv_path.open(encoding="utf-8-sig", newline="") as stream:
        rows = csv.DictReader(stream)
        if rows.fieldnames != ["id", "image", "prompt"]:
            raise ValueError("CSV headers must be exactly id,image,prompt")
        for row in rows:
            job_id = row.get("id") or ""
            if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]{0,63}", job_id) or job_id in ids:
                raise ValueError("Job IDs must be unique ASCII letters, numbers, underscores or hyphens")
            prompt = row.get("prompt") or ""
            if not prompt.strip() or len(prompt) > 4000 or row.get(None):
                raise ValueError("Each row requires a prompt of 1-4000 characters")
            relative = Path(row.get("image") or "")
            if relative.is_absolute() or ".." in relative.parts:
                raise ValueError("Images must be relative paths inside input-root")
            source = (root / relative).resolve(strict=True)
            if not source.is_relative_to(root) or not source.is_file() or source.suffix.lower() != ".png":
                raise ValueError("Only local PNG files inside input-root are supported")
            if source.stat().st_size > MAX_INPUT:
                raise ValueError("PNG exceeds the script's 10 MiB input limit")
            image = source.read_bytes()
            dimensions = png_size(image)
            plan = {"input_sha256": digest(image), "prompt_sha256": digest(prompt.encode()),
                    "model": model, "size": size, "quality": quality, "n": 1}
            plan_hash = digest(json.dumps(plan, sort_keys=True).encode())
            jobs.append({"id": job_id, "source": source, "prompt": prompt, "plan": plan,
                         "plan_hash": plan_hash, "input_dimensions": dimensions})
            ids.add(job_id)
            if len(jobs) > 100:
                raise ValueError("This example supports up to 100 rows")
    if not jobs:
        raise ValueError("CSV requires at least one job")
    return jobs


def multipart(job):
    image = job["source"].read_bytes()
    if len(image) > MAX_INPUT or digest(image) != job["plan"]["input_sha256"]:
        raise ValueError("Input changed after validation; no request sent")
    boundary = "quickrouter-" + uuid.uuid4().hex
    fields = {"model": job["plan"]["model"], "prompt": job["prompt"], "n": "1",
              "size": job["plan"]["size"], "quality": job["plan"]["quality"], "output_format": "png"}
    pieces = []
    for name, value in fields.items():
        pieces.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n'.encode())
    pieces.append(f'--{boundary}\r\nContent-Disposition: form-data; name="image"; filename="input.png"\r\nContent-Type: image/png\r\n\r\n'.encode())
    pieces.extend([image, f"\r\n--{boundary}--\r\n".encode()])
    return b"".join(pieces), "multipart/form-data; boundary=" + boundary


def post_once(body, content_type, key, timeout):
    request = Request(API, data=body, method="POST", headers={
        "Authorization": "Bearer " + key, "Content-Type": content_type, "Accept": "application/json"})
    try:
        with build_opener(NoRedirect()).open(request, timeout=timeout) as response:
            raw = response.read(MAX_RESPONSE + 1)
            if len(raw) > MAX_RESPONSE:
                raise RequestFailure("response_too_large")
            return raw
    except HTTPError as error:
        raise RequestFailure("http_" + str(error.code)) from None
    except (URLError, TimeoutError, OSError):
        raise RequestFailure() from None


def decode_result(raw):
    response = json.loads(raw)
    if not isinstance(response, dict) or response.get("error"):
        raise ValueError("Response is not a successful image object")
    data = response.get("data")
    if not isinstance(data, list) or len(data) != 1 or not isinstance(data[0], dict):
        raise ValueError("Expected exactly one returned image")
    encoded = data[0].get("b64_json")
    if not isinstance(encoded, str):
        raise ValueError("This runner requires b64_json; URLs are not fetched")
    image = base64.b64decode(encoded, validate=True)
    if len(image) > MAX_OUTPUT:
        raise ValueError("Returned image exceeds the script's limit")
    dimensions = png_size(image)
    usage = response.get("usage")
    safe_usage = {}
    if isinstance(usage, dict):
        for field in ("input_tokens", "output_tokens", "total_tokens"):
            if isinstance(usage.get(field), int) and not isinstance(usage[field], bool) and usage[field] >= 0:
                safe_usage[field] = usage[field]
    return image, dimensions, safe_usage


def clean_response(raw):
    image, _, usage = decode_result(raw)
    return json.dumps({"data": [{"b64_json": base64.b64encode(image).decode("ascii")}],
                       "usage": usage}).encode("utf-8")


def recover_result(job, record, out, state, manifest):
    cache = out / (job["id"] + ".response.json")
    if not cache.exists():
        return False
    if cache.is_symlink() or cache.stat().st_size > MAX_RESPONSE:
        raise ValueError("Response cache is invalid")
    try:
        image, dimensions, usage = decode_result(cache.read_bytes())
    except (ValueError, TypeError, KeyError):
        record.update(status="indeterminate", reason="response_not_decodable", updated_at=utc_now())
        save_state(manifest, state)
        return False
    output = out / (job["id"] + ".png")
    private_write(output, image)
    record.update(status="succeeded", reason="file_saved_not_visual_acceptance", updated_at=utc_now(),
                  output=output.name, output_sha256=digest(image), output_dimensions=dimensions,
                  usage=usage, billed_usd=None, visual_acceptance="pending")
    save_state(manifest, state)
    cache.unlink()
    return True


def process_jobs(jobs, out, run=False, key=None, max_requests=1, timeout=180, resume=(),
                 post=post_once):
    manifest = out / "manifest.json"
    state = load_state(manifest)
    selected = set(resume)
    if not selected.issubset({job["id"] for job in jobs}):
        raise ValueError("Resume ID was not found in this CSV")
    for job in jobs:
        record = state["jobs"].get(job["id"])
        if record and record.get("plan_hash") != job["plan_hash"]:
            raise ValueError("Existing job ID has different input, prompt or parameters; keep the original plan")
        if job["id"] in selected and (not record or record.get("status") == "succeeded"):
            raise ValueError("Resume accepts only existing unfinished jobs")
    if not run:
        return {"mode": "offline", "jobs": len(jobs), "new_jobs": sum(job["id"] not in state["jobs"] for job in jobs),
                "existing_jobs": sum(job["id"] in state["jobs"] for job in jobs), "requests_sent": 0,
                "next_step": "Review channel, price, CSV and state before --run"}
    if not key or not key.strip() or "\n" in key or "\r" in key:
        raise ValueError("Set QUICKROUTER_API_KEY in your local environment")
    out.mkdir(parents=True, exist_ok=True, mode=0o700)
    lock = out / ".batch.lock"
    try:
        lock_fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
    except FileExistsError:
        raise ValueError("Another runner or a stale lock exists; inspect the process and manifest first") from None
    sent = saved = skipped = 0
    try:
        with os.fdopen(lock_fd, "w") as stream:
            stream.write(str(os.getpid()))
        state = load_state(manifest)
        for job in jobs:
            record = state["jobs"].get(job["id"])
            if record and record.get("plan_hash") != job["plan_hash"]:
                raise ValueError("Manifest changed; no request sent for this job")
            if record and record.get("status") == "succeeded":
                output = out / (job["id"] + ".png")
                if output.is_symlink() or not output.exists() or digest(output.read_bytes()) != record.get("output_sha256"):
                    raise ValueError("A successful job's file is missing or changed; restore it instead of resubmitting")
                skipped += 1
                continue
            if record and recover_result(job, record, out, state, manifest):
                saved += 1
                continue
            if record and job["id"] not in selected:
                skipped += 1
                continue
            if record and (out / (job["id"] + ".response.json")).exists():
                raise ValueError("Inspect the cached response before permitting another paid request")
            if not record and ((out / (job["id"] + ".png")).exists() or
                               (out / (job["id"] + ".response.json")).exists()):
                raise ValueError("An untracked output exists; restore the manifest before submitting")
            if sent >= max_requests:
                break
            body, content_type = multipart(job)
            record = {"plan_hash": job["plan_hash"], "plan": job["plan"],
                      "input_dimensions": job["input_dimensions"], "status": "indeterminate",
                      "reason": "request_in_flight", "updated_at": utc_now(),
                      "attempts": (record or {}).get("attempts", 0) + 1}
            state["jobs"][job["id"]] = record
            save_state(manifest, state)
            sent += 1
            try:
                raw = post(body, content_type, key, timeout)
            except RequestFailure as failure:
                record.update(reason=failure.code, updated_at=utc_now())
                save_state(manifest, state)
                break
            try:
                cleaned = clean_response(raw)
            except (ValueError, TypeError, KeyError):
                record.update(reason="response_not_decodable", updated_at=utc_now())
                save_state(manifest, state)
                break
            private_write(out / (job["id"] + ".response.json"), cleaned)
            record.update(status="response_received", reason="saved_response", updated_at=utc_now())
            save_state(manifest, state)
            if not recover_result(job, record, out, state, manifest):
                break
            saved += 1
        return {"mode": "run", "jobs": len(jobs), "requests_sent": sent, "files_saved": saved,
                "skipped": skipped, "indeterminate_jobs": sum(
                    state["jobs"].get(job["id"], {}).get("status") == "indeterminate" for job in jobs),
                "manifest": "manifest.json", "billing_verified": False}
    finally:
        lock.unlink()


def main():
    parser = argparse.ArgumentParser(description="Serial PNG edits; default offline; no automatic POST retries")
    parser.add_argument("--jobs", type=Path, required=True)
    parser.add_argument("--input-root", type=Path, required=True)
    parser.add_argument("--out", type=Path, default=Path("batch-output"))
    parser.add_argument("--model", default="gpt-image-2", choices=["gpt-image-2"])
    parser.add_argument("--size", default="1024x1024", choices=["1024x1024", "1024x1536", "1536x1024", "auto"])
    parser.add_argument("--quality", default="low", choices=["low", "medium", "high", "auto"])
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--max-requests", type=int, default=1)
    parser.add_argument("--timeout", type=int, default=180)
    parser.add_argument("--resume-job", action="append", default=[])
    parser.add_argument("--acknowledge-new-charge", action="store_true")
    args = parser.parse_args()
    if not 1 <= args.max_requests <= 100 or not 30 <= args.timeout <= 600:
        parser.error("max-requests must be 1-100; timeout must be 30-600 seconds")
    if args.resume_job and not (args.run and args.acknowledge_new_charge):
        parser.error("Resubmission requires --run --acknowledge-new-charge after console verification")
    try:
        jobs = load_jobs(args.jobs, args.input_root, args.model, args.size, args.quality)
        result = process_jobs(jobs, args.out.resolve(), run=args.run,
                              key=os.environ.get("QUICKROUTER_API_KEY") if args.run else None,
                              max_requests=args.max_requests, timeout=args.timeout, resume=args.resume_job)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        if result.get("indeterminate_jobs", 0):
            return 1
    except (ValueError, OSError, csv.Error, json.JSONDecodeError):
        print("Stopped safely. Check CSV, local files, manifest or cached response. No automatic retry; credentials and API bodies are not printed.", file=sys.stderr)
        return 1
    return 0


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