"""HamuBox NovelAI native stream example. Python 3.10+; pip install msgpack.

Set HAMUBOX_API_KEY, then:
  python novelai-stream.py novelai-native-stream.json --output result.png
No automatic POST retry. Resume status using the saved task ID and client.py.
"""
import argparse
import json
import os
from pathlib import Path
import struct
import sys
import urllib.error
import urllib.request

import msgpack

BASE_URL = "https://api.hamubox.com/v1/novelai"
MAX_FRAME_BYTES = 32 * 1024 * 1024


def read_exact(response, size):
    data = bytearray()
    while len(data) < size:
        chunk = response.read(size - len(data))
        if not chunk:
            raise EOFError("Incomplete stream; query the saved task ID before retrying generation.")
        data.extend(chunk)
    return bytes(data)


def read_final(response):
    while True:
        size = struct.unpack(">I", read_exact(response, 4))[0]
        if not 0 < size <= MAX_FRAME_BYTES:
            raise RuntimeError("Invalid stream frame length.")
        event = msgpack.unpackb(read_exact(response, size), raw=False)
        if not isinstance(event, dict):
            raise RuntimeError("Invalid stream event.")
        kind = event.get("event_type")
        if kind == "error":
            raise RuntimeError(event.get("message") or "Generation stream failed.")
        if kind == "intermediate":
            print(f"Step {event.get('step_ix')}", file=sys.stderr)
        if kind == "final":
            image = event.get("image")
            if event.get("samp_ix", 0) != 0 or not isinstance(image, bytes) or not image:
                raise RuntimeError("Invalid final image.")
            # n_samples=1: a complete final is enough; do not wait for clean EOF.
            return image


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("request_file")
    parser.add_argument("--output", default="result.png")
    args = parser.parse_args()
    key = os.environ.get("HAMUBOX_API_KEY", "").strip()
    if not key:
        parser.error("Set HAMUBOX_API_KEY first.")
    with open(args.request_file, encoding="utf-8") as source:
        payload = json.load(source)
    if "payload" in payload or payload.get("parameters", {}).get("n_samples", 1) != 1:
        parser.error("Use native generation JSON with n_samples=1, not an asynchronous task wrapper.")
    req = urllib.request.Request(
        BASE_URL + "/ai/generate-image-stream", method="POST",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Authorization": "Bearer " + key, "Content-Type": "application/json"},
    )
    # Never automatically retry: disconnecting does not cancel the accepted task.
    with urllib.request.urlopen(req, timeout=660) as response:
        task_id = response.headers.get("X-NovelAI-Task-Id")
        if not task_id:
            raise RuntimeError("Missing task ID; do not automatically retry POST.")
        print(f"TASK_ID={task_id}", flush=True)
        output = Path(args.output)
        output.with_suffix(".task.json").write_text(
            json.dumps({"id": task_id, "kind": "novelai", "model": payload.get("model")}),
            encoding="utf-8",
        )
        if response.headers.get_content_type() != "application/msgpack":
            raise RuntimeError("Unexpected content type; query the saved task ID.")
        image = read_final(response)
    output.write_bytes(image)
    print(f"Image saved to {output}; storage and settlement may still be in progress.")
    print(f"python client.py novelai --task-id {task_id}")


if __name__ == "__main__":
    try:
        main()
    except urllib.error.HTTPError as error:
        print(f"HTTP {error.code}: " + error.read(4096).decode("utf-8", errors="replace"), file=sys.stderr)
        sys.exit(1)
    except (OSError, EOFError, RuntimeError, ValueError, urllib.error.URLError) as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)