from __future__ import annotations

import argparse
import base64
import json
import struct
import sys
import urllib.error
import urllib.request
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_image_data_url(width: int = 512, height: int = 512) -> str:
    pixels = []
    for y in range(height):
        for x in range(width):
            if x < width // 2 and y < height // 2:
                pixels.append((230, 70, 70))
            elif x >= width // 2 and y < height // 2:
                pixels.append((70, 180, 95))
            elif x < width // 2:
                pixels.append((70, 100, 220))
            else:
                pixels.append((240, 210, 70))
    return "data:image/png;base64," + base64.b64encode(rgb_png(width, height, pixels)).decode("ascii")


def post_json(url: str, payload: dict, timeout: int) -> 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}")
    return data


def run(args: argparse.Namespace) -> int:
    payload = {
        "base_url": args.video_base_url,
        "image_base64": make_image_data_url(args.width, args.height),
        "filename": "smoke-grid.png",
        "prompt": args.prompt,
        "model": args.model,
        "max_tokens": args.max_tokens,
    }
    result = post_json(f"{args.chat_base_url.rstrip('/')}/api/video/analyze", payload, args.timeout)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if result.get("error"):
        return 1
    return 0 if result.get("text") and result.get("model") else 1


def format_url_error(exc: urllib.error.URLError) -> str:
    if isinstance(exc, urllib.error.HTTPError):
        try:
            body = exc.read().decode("utf-8", errors="replace").strip()
        except OSError:
            body = ""
        return f"HTTP {exc.code}: {body or exc.reason}"
    return str(exc)


def main() -> int:
    parser = argparse.ArgumentParser(description="Smoke test vision analysis through the model chat API.")
    parser.add_argument("--chat-base-url", default="http://192.168.220.91:8765")
    parser.add_argument("--video-base-url", default="http://docker-gpu.cin.su:8030")
    parser.add_argument("--model", default="qwen2.5-vl-7b-instruct")
    parser.add_argument("--prompt", default="Опиши изображение: какие крупные цветные области видны?")
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--allow-unavailable", action="store_true")
    args = parser.parse_args()

    try:
        return run(args)
    except (OSError, ValueError, urllib.error.URLError) as exc:
        message = format_url_error(exc) if isinstance(exc, urllib.error.URLError) else str(exc)
        print(f"Smoke video analysis failed: {message}", file=sys.stderr)
        return 0 if args.allow_unavailable else 1


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