Genarrative/.codex/skills/genarrative-external-editor-api/scripts/genarrative_external_api.py

#!/usr/bin/env python3
"""Tiny stdlib client for Genarrative external editor APIs."""

from __future__ import annotations

import argparse
import json
import mimetypes
import os
import re
import struct
import sys
import tempfile
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path
from typing import Any


BASE_URL = "https://www.genarrative.world/"
DEFAULT_CREDENTIALS_FILE = Path.home() / ".config/genarrative/external-editor-api.json"


class GenarrativeApiError(RuntimeError):
    pass


def load_api_key(credentials_file: str | os.PathLike[str] = DEFAULT_CREDENTIALS_FILE) -> str:
    path = Path(credentials_file).expanduser()
    try:
        value = json.loads(path.read_text(encoding="utf-8")).get("apiKey", "")
    except FileNotFoundError as error:
        raise GenarrativeApiError(
            f"Missing API key file: {path}. Create JSON like {{\"apiKey\":\"tnr_sk_...\"}}."
        ) from error
    except json.JSONDecodeError as error:
        raise GenarrativeApiError(f"Invalid JSON in API key file: {path}.") from error
    if not isinstance(value, str) or not value.strip():
        raise GenarrativeApiError(f"Missing apiKey string in API key file: {path}.")
    return value.strip()


def unwrap_envelope(body: Any) -> Any:
    if isinstance(body, dict) and body.get("ok") is True and "data" in body:
        return body["data"]
    return body


def guess_content_type(file_path: str | os.PathLike[str]) -> str:
    return mimetypes.guess_type(str(file_path))[0] or "application/octet-stream"


def source_layer_id_from_path(file_path: str | os.PathLike[str]) -> str:
    stem = Path(file_path).stem.lower()
    slug = re.sub(r"[^a-z0-9]+", "-", stem).strip("-") or "image"
    return f"external-reference-{slug}"


def image_dimensions(file_path: str | os.PathLike[str]) -> tuple[int, int] | None:
    path = Path(file_path)
    with path.open("rb") as fh:
        header = fh.read(32)
        if header.startswith(b"\x89PNG\r\n\x1a\n") and header[12:16] == b"IHDR":
            return struct.unpack(">II", header[16:24])
        if not header.startswith(b"\xff\xd8"):
            return None
        fh.seek(2)
        while True:
            marker_prefix = fh.read(1)
            if not marker_prefix:
                return None
            if marker_prefix != b"\xff":
                continue
            marker = fh.read(1)
            while marker == b"\xff":
                marker = fh.read(1)
            if marker in {b"\xd8", b"\xd9"}:
                continue
            length_bytes = fh.read(2)
            if len(length_bytes) != 2:
                return None
            length = struct.unpack(">H", length_bytes)[0]
            if marker and marker[0] in {0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF}:
                data = fh.read(5)
                if len(data) != 5:
                    return None
                height, width = struct.unpack(">HH", data[1:5])
                return width, height
            fh.seek(max(length - 2, 0), os.SEEK_CUR)


class GenarrativeExternalClient:
    def __init__(
        self,
        api_key: str | None = None,
        credentials_file: str | os.PathLike[str] = DEFAULT_CREDENTIALS_FILE,
        base_url: str = BASE_URL,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key if api_key is not None else load_api_key(credentials_file)

    def request_json(
        self,
        method: str,
        path: str,
        body: dict[str, Any] | None = None,
        query: dict[str, Any] | None = None,
        auth: bool = True,
        timeout: int = 60,
    ) -> Any:
        url = f"{self.base_url}{path}"
        if query:
            url = f"{url}?{urllib.parse.urlencode({k: v for k, v in query.items() if v is not None})}"
        data = None if body is None else json.dumps(body).encode("utf-8")
        headers = {"Accept": "application/json"}
        if data is not None:
            headers["Content-Type"] = "application/json"
        if auth:
            headers["Authorization"] = f"Bearer {self.api_key}"
        request = urllib.request.Request(url, data=data, headers=headers, method=method.upper())
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                payload = response.read()
        except urllib.error.HTTPError as error:
            detail = error.read().decode("utf-8", errors="replace")
            raise GenarrativeApiError(f"{method.upper()} {path} failed: HTTP {error.code}: {detail}") from error
        if not payload:
            return None
        return unwrap_envelope(json.loads(payload.decode("utf-8")))

    def openapi(self) -> Any:
        return self.request_json("GET", "/api/external/v1/openapi.json", auth=False)

    def list_projects(self) -> Any:
        return self.request_json("GET", "/api/external/v1/editor/projects")

    def create_project(self, title: str | None = None) -> Any:
        body = {} if title is None else {"title": title}
        return self.request_json("POST", "/api/external/v1/editor/projects", body)

    def save_canvas(self, project_id: str, viewport: dict[str, Any], layers: dict[str, Any]) -> Any:
        return self.request_json(
            "PATCH",
            f"/api/external/v1/editor/projects/{urllib.parse.quote(project_id, safe='')}/canvas",
            {"viewport": viewport, "layers": layers},
        )

    def create_upload_ticket(self, file_path: str | os.PathLike[str], access: str = "private") -> Any:
        path = Path(file_path)
        return self.request_json(
            "POST",
            "/api/external/v1/assets/direct-upload-tickets",
            {
                "legacyPrefix": "generated-character-drafts",
                "pathSegments": ["editor", "external-editor-references"],
                "fileName": path.name,
                "contentType": guess_content_type(path),
                "access": access,
            },
        )

    def upload_to_oss(self, ticket_response: Any, file_path: str | os.PathLike[str]) -> None:
        payload = unwrap_envelope(ticket_response)
        ticket = payload.get("upload") if isinstance(payload, dict) else None
        if not isinstance(ticket, dict):
            raise GenarrativeApiError("Upload ticket response missing upload payload.")
        fields = ticket.get("formFields")
        if not isinstance(fields, dict):
            raise GenarrativeApiError("Upload ticket response missing formFields.")
        boundary = f"----genarrative-{uuid.uuid4().hex}"
        path = Path(file_path)
        content_type = ticket.get("contentType") or guess_content_type(path)
        chunks: list[bytes] = []
        for key, value in fields.items():
            if value is None:
                continue
            chunks.extend(
                [
                    f"--{boundary}\r\n".encode(),
                    f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode(),
                    str(value).encode(),
                    b"\r\n",
                ]
            )
        chunks.extend(
            [
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="file"; filename="{path.name}"\r\n'.encode(),
                f"Content-Type: {content_type}\r\n\r\n".encode(),
                path.read_bytes(),
                b"\r\n",
                f"--{boundary}--\r\n".encode(),
            ]
        )
        request = urllib.request.Request(
            ticket["host"],
            data=b"".join(chunks),
            headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=300) as response:
                if response.status not in {200, 201, 204}:
                    raise GenarrativeApiError(f"OSS upload failed: HTTP {response.status}")
        except urllib.error.HTTPError as error:
            detail = error.read().decode("utf-8", errors="replace")
            raise GenarrativeApiError(f"OSS upload failed: HTTP {error.code}: {detail}") from error

    def confirm_asset_object(
        self,
        object_key: str,
        file_path: str | os.PathLike[str],
        asset_kind: str = "editor_reference_image",
        access_policy: str = "private",
    ) -> Any:
        path = Path(file_path)
        return self.request_json(
            "POST",
            "/api/external/v1/assets/objects/confirm",
            {
                "objectKey": object_key,
                "contentType": guess_content_type(path),
                "contentLength": path.stat().st_size,
                "assetKind": asset_kind,
                "accessPolicy": access_policy,
            },
        )

    def upload_reference_image(self, file_path: str | os.PathLike[str]) -> dict[str, Any]:
        ticket = self.create_upload_ticket(file_path)
        upload = unwrap_envelope(ticket)["upload"]
        self.upload_to_oss(ticket, file_path)
        confirmed = self.confirm_asset_object(upload["objectKey"], file_path)
        return {
            "objectKey": upload["objectKey"],
            "ticket": ticket,
            "assetObject": unwrap_envelope(confirmed).get("assetObject"),
            "dimensions": image_dimensions(file_path),
            "sourceLayerId": source_layer_id_from_path(file_path),
        }

    def read_url(self, object_key: str) -> Any:
        return self.request_json("GET", "/api/external/v1/assets/read-url", query={"objectKey": object_key})

    def generate_image(self, prompt: str, **fields: Any) -> Any:
        return self.request_json("POST", "/api/external/v1/editor/images/generations", {"prompt": prompt, **fields})

    def edit_image(self, prompt: str, source_image_src: str, **fields: Any) -> Any:
        return self.request_json(
            "POST",
            "/api/external/v1/editor/images/edits",
            {"prompt": prompt, "sourceImageSrc": source_image_src, **fields},
        )

    def extract_ui_assets(self, source_image_src: str, image_size: str = "1K", **fields: Any) -> Any:
        fields.pop("aspectRatio", None)
        return self.request_json(
            "POST",
            "/api/external/v1/editor/ui-designs/assets/extractions",
            {"sourceImageSrc": source_image_src, "imageSize": image_size, **fields, "aspectRatio": "1:1"},
        )

    def animate_character(
        self,
        source_image_src: str,
        source_width: int,
        source_height: int,
        prompt_text: str,
        source_layer_id: str,
        **fields: Any,
    ) -> Any:
        body = {
            "sourceLayerId": source_layer_id,
            "sourceImageSrc": source_image_src,
            "sourceWidth": source_width,
            "sourceHeight": source_height,
            "promptText": prompt_text,
            "resolution": fields.pop("resolution", "720p"),
            "ratio": fields.pop("ratio", "same"),
            "frameCount": fields.pop("frameCount", 40),
            "durationSeconds": fields.pop("durationSeconds", 5),
            **fields,
            "model": "seedance2.0-fast",
        }
        return self.request_json("POST", "/api/external/v1/editor/character-animations/generations", body)

    def generate_video(self, prompt: str, **fields: Any) -> Any:
        fields.pop("mode", None)
        body = {
            "prompt": prompt,
            "model": fields.pop("model", "seedance2.0-fast"),
            "aspectRatio": fields.pop("aspectRatio", "16:9"),
            "durationSeconds": fields.pop("durationSeconds", 5),
            "resolution": fields.pop("resolution", "720p"),
            "sound": fields.pop("sound", "off"),
            **fields,
            "mode": "std",
        }
        return self.request_json("POST", "/api/external/v1/editor/videos/generations", body)

    def generate_sound_effect(self, prompt: str, duration: int, **fields: Any) -> Any:
        return self.request_json(
            "POST",
            "/api/external/v1/editor/audios/sound-effects/generations",
            {"prompt": prompt, "duration": duration, **fields},
        )

    def generate_background_music(self, description: str, **fields: Any) -> Any:
        return self.request_json(
            "POST",
            "/api/external/v1/editor/audios/background-music/generations",
            {"gptDescriptionPrompt": description, **fields, "makeInstrumental": True},
        )


def _self_test() -> None:
    png = (
        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"
        b"\x00\x00\x00\x02\x00\x00\x00\x03\x08\x06\x00\x00\x00"
    )
    with tempfile.NamedTemporaryFile(suffix="Hero Image.png") as fh:
        fh.write(png)
        fh.flush()
        assert image_dimensions(fh.name) == (2, 3)
        assert source_layer_id_from_path(fh.name).startswith("external-reference-")
    assert unwrap_envelope({"ok": True, "data": {"upload": 1}}) == {"upload": 1}
    print("self-test ok")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--credentials-file", default=str(DEFAULT_CREDENTIALS_FILE))
    parser.add_argument("command", nargs="?", choices=["openapi", "list-projects"])
    args = parser.parse_args(argv)
    if args.self_test:
        _self_test()
        return 0
    if args.command is None:
        parser.print_help()
        return 0
    if args.command == "openapi":
        client = GenarrativeExternalClient(api_key="", credentials_file=args.credentials_file)
        print(json.dumps(client.openapi(), ensure_ascii=False, indent=2))
    elif args.command == "list-projects":
        client = GenarrativeExternalClient(credentials_file=args.credentials_file)
        print(json.dumps(client.list_projects(), ensure_ascii=False, indent=2))
    return 0


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