from __future__ import annotations

import argparse
import base64
import json
import struct
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
import zlib


def png_chunk(kind: bytes, data: bytes) -> bytes:
    return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)


def rgb_png(width: int, height: int, pixels: list[tuple[int, int, int]]) -> bytes:
    rows = []
    for y in range(height):
        start = y * width
        row = b"".join(bytes(pixel) for pixel in pixels[start : start + width])
        rows.append(b"\x00" + row)
    header = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    return b"\x89PNG\r\n\x1a\n" + png_chunk(b"IHDR", header) + png_chunk(b"IDAT", zlib.compress(b"".join(rows), 6)) + png_chunk(b"IEND", b"")


def make_source_image(width: int, height: int) -> str:
    pixels = []
    for y in range(height):
        for x in range(width):
            if width // 3 < x < 2 * width // 3 and height // 3 < y < 2 * height // 3:
                pixels.append((210, 60, 70))
            else:
                pixels.append((35 + (x * 80 // width), 95 + (y * 80 // height), 145))
    return "data:image/png;base64," + base64.b64encode(rgb_png(width, height, pixels)).decode("ascii")


def make_mask_image(width: int, height: int) -> str:
    pixels = []
    for y in range(height):
        for x in range(width):
            masked = width // 3 < x < 2 * width // 3 and height // 3 < y < 2 * height // 3
            pixels.append((255, 255, 255) if masked else (0, 0, 0))
    return "data:image/png;base64," + base64.b64encode(rgb_png(width, height, pixels)).decode("ascii")


def post_json(url: str, payload: dict, timeout: int = 30) -> dict:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        data = json.loads(response.read().decode("utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"JSON object expected from {url}")
    if data.get("error"):
        raise ValueError(str(data["error"]))
    return data


def get_json(url: str, timeout: int = 30) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as response:
        data = json.loads(response.read().decode("utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"JSON object expected from {url}")
    if data.get("error"):
        raise ValueError(str(data["error"]))
    return data


def submit_job(args: argparse.Namespace, operation: str) -> dict:
    payload: dict[str, object] = {
        "operation": operation,
        "base_url": args.image_base_url,
        "model_id": args.model_id,
        "model_mode": args.model_mode,
        "model": args.model,
        "prompt": args.prompt,
        "negative_prompt": args.negative_prompt,
        "width": args.width,
        "height": args.height,
        "steps": args.steps,
        "guidance_scale": args.guidance_scale,
        "seed": args.seed,
    }
    if operation == "edit":
        payload.update(
            {
                "image_base64": make_source_image(args.width, args.height),
                "mask_base64": make_mask_image(args.width, args.height),
                "strength": args.strength,
            }
        )
    return post_json(f"{args.chat_base_url.rstrip('/')}/api/image/submit", payload)


def poll_job(args: argparse.Namespace, job_id: str, base_url: str) -> dict:
    deadline = time.monotonic() + args.timeout
    last_status = ""
    while time.monotonic() < deadline:
        url = (
            f"{args.chat_base_url.rstrip('/')}/api/image/job?"
            f"job_id={urllib.parse.quote(job_id)}&base_url={urllib.parse.quote(base_url, safe='')}"
        )
        job = get_json(url)
        status = str(job.get("status") or "")
        if status != last_status:
            print(f"{job_id}: {status}")
            last_status = status
        if status in {"completed", "cancelled", "error"}:
            return job
        time.sleep(args.poll_interval)
    raise TimeoutError(f"image job {job_id} did not finish within {args.timeout}s")


def cancel_job(args: argparse.Namespace, job_id: str, base_url: str) -> dict:
    return post_json(
        f"{args.chat_base_url.rstrip('/')}/api/image/cancel",
        {"job_id": job_id, "base_url": base_url},
    )


def run_one(args: argparse.Namespace, operation: str) -> int:
    submitted = submit_job(args, operation)
    job_id = str(submitted.get("report_id") or submitted.get("id") or uuid.uuid4())
    base_url = str(submitted.get("base_url") or args.image_base_url)
    print(f"submitted {operation}: {job_id}")
    if args.cancel:
        time.sleep(args.cancel_after)
        cancel_job(args, job_id, base_url)
    job = poll_job(args, job_id, base_url)
    status = job.get("status")
    print(json.dumps({key: job.get(key) for key in ("id", "operation", "status", "latency_ms", "artifact", "error")}, ensure_ascii=False, indent=2))
    if args.cancel:
        return 0 if status in {"cancelled", "cancel_requested"} else 1
    return 0 if status == "completed" and job.get("artifact") else 1


def main() -> int:
    parser = argparse.ArgumentParser(description="Smoke test async image jobs through the model chat API.")
    parser.add_argument("--chat-base-url", default="http://192.168.220.91:8765")
    parser.add_argument("--image-base-url", default="http://docker-gpu.cin.su:8040")
    parser.add_argument("--model", default="sdxl-image")
    parser.add_argument("--model-id", default="sdxl-base-1_0")
    parser.add_argument("--model-mode", default="sdxl")
    parser.add_argument("--operation", choices=["generate", "edit"], default="generate")
    parser.add_argument("--prompt", default="A small clean product photo of a blue ceramic mug on a white desk")
    parser.add_argument("--negative-prompt", default="blurry, distorted, low quality")
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--steps", type=int, default=8)
    parser.add_argument("--guidance-scale", type=float, default=5.5)
    parser.add_argument("--strength", type=float, default=0.85)
    parser.add_argument("--seed", type=int, default=123)
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--poll-interval", type=float, default=2.0)
    parser.add_argument("--cancel", action="store_true")
    parser.add_argument("--cancel-after", type=float, default=1.0)
    args = parser.parse_args()

    try:
        return run_one(args, args.operation)
    except (OSError, TimeoutError, ValueError, urllib.error.URLError) as exc:
        print(f"Smoke image job failed: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
