import argparse
import base64
import hashlib
import json
import os
from pathlib import Path
import struct
import sys


def main():
  parser = argparse.ArgumentParser(description="QuickRouter Images SDK sample; dry-run unless --run")
  parser.add_argument("--run", action="store_true")
  parser.add_argument("--mode", choices=["generate", "edit"], default="generate")
  parser.add_argument("--model", default="gpt-image-2")
  parser.add_argument("--prompt-file", default="prompt.txt")
  parser.add_argument("--image")
  parser.add_argument("--output", default="image-result.png")
  args = parser.parse_args()
  prompt_path = Path(args.prompt_file)
  if not prompt_path.is_file() or prompt_path.stat().st_size > 32000:
    raise ValueError("prompt file missing or exceeds 32 KB")
  prompt = prompt_path.read_text(encoding="utf-8").strip()
  if not prompt or len(args.model) > 128:
    raise ValueError("prompt/model invalid")
  output = Path(args.output)
  metadata = output.with_suffix(output.suffix + ".json")
  if not output.parent.is_dir():
    raise ValueError("output parent directory must already exist")
  if not os.access(output.parent, os.W_OK):
    raise ValueError("output parent directory is not writable")
  if output.exists() or metadata.exists() or output.is_symlink() or metadata.is_symlink():
    raise ValueError("output already exists; choose a new filename")
  image = Path(args.image) if args.image else None
  if args.mode == "generate" and image:
    raise ValueError("--image is only valid with --mode edit")
  if args.mode == "edit" and (not image or not image.is_file() or image.stat().st_size > 20000000):
    raise ValueError("edit requires a local image file <=20 MB")
  params = {"model": args.model, "prompt": prompt, "n": 1, "size": "1024x1024", "output_format": "png"}
  if not args.run:
    print(json.dumps({"dry_run": True, "endpoint": "/v1/images/edits" if image else "/v1/images/generations", "model": args.model, "size": params["size"], "image_attached": bool(image), "output": str(output), "note": "No API request sent. Use --run only after checking channel support and price."}, ensure_ascii=False))
    return
  key = os.environ.get("QUICKROUTER_API_KEY", "").strip()
  if not key:
    raise ValueError("QUICKROUTER_API_KEY missing")
  from openai import OpenAI
  client = OpenAI(api_key=key, base_url="https://api.quickrouter.ai/v1", max_retries=0, timeout=180.0)
  try:
    if image:
      with image.open("rb") as file:
        raw = client.images.with_raw_response.edit(image=file, **params)
    else:
      raw = client.images.with_raw_response.generate(**params)
    result = raw.parse()
  except Exception as error:
    print(json.dumps({"status": "indeterminate_or_error", "type": type(error).__name__, "http_status": getattr(error, "status_code", None), "request_id": getattr(error, "request_id", None), "note": "No automatic resend. Check request logs and charges before another attempt."}), file=sys.stderr)
    raise SystemExit(1)
  request_id = raw.headers.get("x-request-id")
  encoded = result.data[0].b64_json if result.data else None
  if not encoded:
    raise ValueError("No b64_json. URL responses are not automatically fetched; review the channel protocol and download manually without API authorization headers.")
  if len(encoded) > 45000000:
    raise ValueError("image response exceeds local 32 MB limit; do not resubmit automatically")
  image_bytes = base64.b64decode(encoded, validate=True)
  if not image_bytes.startswith(b"\x89PNG\r\n\x1a\n") or len(image_bytes) < 24 or len(image_bytes) > 32000000:
    raise ValueError("response is not a valid-size PNG; do not resubmit automatically")
  width, height = struct.unpack(">II", image_bytes[16:24])
  usage = getattr(result, "usage", None)
  record = {"request_id": request_id, "model": args.model, "mode": args.mode, "requested_size": "1024x1024", "actual_width": width, "actual_height": height, "sha256": hashlib.sha256(image_bytes).hexdigest(), "usage": usage.model_dump() if usage else None, "billing_status": "unverified", "actual_charge": None}
  with output.open("xb") as file:
    file.write(image_bytes)
  with metadata.open("x", encoding="utf-8") as file:
    json.dump(record, file, ensure_ascii=False, indent=2)
  print(json.dumps({"saved": str(output), "metadata": str(metadata), "request_id": request_id, "billing_status": "unverified"}, ensure_ascii=False))


if __name__ == "__main__":
  main()