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"
DEFAULT_REQUEST_TIMEOUT_SECONDS = 60
GENERATION_REQUEST_TIMEOUT_SECONDS = 420


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 normalize_optional_text(value: Any) -> str | None:
    if not isinstance(value, str):
        return None
    stripped = value.strip()
    return stripped or None


def art_spec_prompt(prompt: str, art_spec: dict[str, Any] | None) -> str:
    if not art_spec:
        return prompt
    spec_json = json.dumps(art_spec, ensure_ascii=False, sort_keys=True)
    return f"{prompt}\n\n美术规范(JSON): {spec_json}"


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 = DEFAULT_REQUEST_TIMEOUT_SECONDS,
    ) -> 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 list_asset_library(self) -> Any:
        return self.request_json("GET", "/api/external/v1/editor/assets/library")

    def create_asset_folder(self, label: str, sort_order: int = 100) -> Any:
        return self.request_json(
            "POST",
            "/api/external/v1/editor/assets/folders",
            {"label": label, "sortOrder": sort_order},
        )

    def ensure_asset_folder(self, label: str, sort_order: int = 100) -> dict[str, Any]:
        normalized_label = normalize_optional_text(label) or "新画板"
        library = self.list_asset_library()
        folders = unwrap_envelope(library).get("library", {}).get("folders", [])
        if isinstance(folders, list):
            for folder in folders:
                if isinstance(folder, dict) and normalize_optional_text(folder.get("label")) == normalized_label:
                    return folder
        created = self.create_asset_folder(normalized_label, sort_order=sort_order)
        folder = unwrap_envelope(created).get("folder")
        if not isinstance(folder, dict):
            raise GenarrativeApiError("Create asset folder response missing folder payload.")
        return folder

    def create_asset(
        self,
        folder_id: str,
        label: str,
        image_src: str,
        width: int,
        height: int,
        **fields: Any,
    ) -> Any:
        body = {
            "folderId": folder_id,
            "label": label,
            "imageSrc": image_src,
            "width": width,
            "height": height,
            "sourceType": fields.pop("sourceType", "generated"),
            **fields,
        }
        return self.request_json("POST", "/api/external/v1/editor/assets", body)

    def prepare_canvas_session(self, canvas_name: str) -> dict[str, Any]:
        normalized_name = normalize_optional_text(canvas_name) or "新画板"
        project = unwrap_envelope(self.create_project(normalized_name)).get("project")
        if not isinstance(project, dict) or not normalize_optional_text(project.get("projectId")):
            raise GenarrativeApiError("Create project response missing projectId.")
        folder = self.ensure_asset_folder(normalized_name)
        folder_id = normalize_optional_text(folder.get("folderId"))
        if not folder_id:
            raise GenarrativeApiError("Asset folder payload missing folderId.")
        return {
            "canvasName": normalized_name,
            "projectId": project["projectId"],
            "assetFolderId": folder_id,
            "project": project,
            "folder": folder,
        }

    def build_canvas_completion(
        self,
        title: str,
        width: int,
        height: int,
        x: float = 0,
        y: float = 0,
        dialog_id: str | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "title": normalize_optional_text(title) or "生成素材",
            "placeholder": {
                "x": x,
                "y": y,
                "width": width,
                "height": height,
                "originalWidth": width,
                "originalHeight": height,
            },
        }
        if normalize_optional_text(dialog_id):
            payload["dialogId"] = dialog_id
        return payload

    def canvas_generation_fields(
        self,
        session: dict[str, Any],
        asset_label: str,
        width: int = 1024,
        height: int = 1024,
        x: float = 0,
        y: float = 0,
        dialog_id: str | None = None,
        asset_label_field: str | None = "assetLabel",
    ) -> dict[str, Any]:
        fields: dict[str, Any] = {
            "projectId": session["projectId"],
            "canvasCompletion": self.build_canvas_completion(
                asset_label,
                width=width,
                height=height,
                x=x,
                y=y,
                dialog_id=dialog_id,
            ),
        }
        folder_id = normalize_optional_text(session.get("assetFolderId"))
        if folder_id:
            fields["assetFolderId"] = folder_id
        if asset_label_field:
            fields[asset_label_field] = normalize_optional_text(asset_label) or "生成素材"
        return fields

    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 _apply_art_spec(self, fields: dict[str, Any], prompt: str) -> str:
        art_spec = fields.pop("artSpec", None)
        if art_spec is None:
            art_spec = fields.pop("art_spec", None)
        if not isinstance(art_spec, dict):
            return prompt
        generation_inputs = fields.get("generationInputs")
        if not isinstance(generation_inputs, dict):
            generation_inputs = {}
        generation_inputs.setdefault("artSpec", art_spec)
        fields["generationInputs"] = generation_inputs
        return art_spec_prompt(prompt, art_spec)

    def _apply_canvas_session_fields(
        self,
        fields: dict[str, Any],
        default_label: str,
        default_width: int,
        default_height: int,
        asset_label_field: str | None = "assetLabel",
    ) -> tuple[dict[str, Any] | None, str]:
        session = fields.pop("canvasSession", None)
        if session is None:
            session = fields.pop("canvas_session", None)
        if isinstance(session, str):
            session = self.prepare_canvas_session(session)
        label = (
            normalize_optional_text(fields.get(asset_label_field)) if asset_label_field else None
        ) or normalize_optional_text(fields.pop("canvasTitle", None)) or normalize_optional_text(default_label) or "生成素材"
        if asset_label_field and not normalize_optional_text(fields.get(asset_label_field)):
            fields.pop(asset_label_field, None)
        if not isinstance(session, dict):
            return None, label
        width = int(fields.pop("canvasWidth", default_width))
        height = int(fields.pop("canvasHeight", default_height))
        x = float(fields.pop("canvasX", 0))
        y = float(fields.pop("canvasY", 0))
        dialog_id = normalize_optional_text(fields.pop("dialogId", None))
        fields.update(
            {
                key: value
                for key, value in self.canvas_generation_fields(
                    session,
                    label,
                    width=width,
                    height=height,
                    x=x,
                    y=y,
                    dialog_id=dialog_id,
                    asset_label_field=asset_label_field,
                ).items()
                if key not in fields or fields[key] is None
            }
        )
        return session, label

    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:
        self._apply_canvas_session_fields(fields, prompt, 1024, 1024)
        prompt = self._apply_art_spec(fields, prompt)
        return self.request_json(
            "POST",
            "/api/external/v1/editor/images/generations",
            {"prompt": prompt, **fields},
            timeout=GENERATION_REQUEST_TIMEOUT_SECONDS,
        )

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

    def extract_ui_assets(self, source_image_src: str, image_size: str = "1K", **fields: Any) -> Any:
        fields.pop("aspectRatio", None)
        self._apply_canvas_session_fields(fields, fields.get("spritesheetLabel", "UI 素材拆分"), 1024, 1024, "spritesheetLabel")
        return self.request_json(
            "POST",
            "/api/external/v1/editor/ui-designs/assets/extractions",
            {"sourceImageSrc": source_image_src, "imageSize": image_size, **fields, "aspectRatio": "1:1"},
            timeout=GENERATION_REQUEST_TIMEOUT_SECONDS,
        )

    def animate_character(
        self,
        source_image_src: str,
        source_width: int,
        source_height: int,
        prompt_text: str,
        source_layer_id: str,
        **fields: Any,
    ) -> Any:
        session, asset_label = self._apply_canvas_session_fields(
            fields,
            fields.get("canvasTitle", "角色动画"),
            source_width,
            source_height,
            asset_label_field=None,
        )
        prompt_text = self._apply_art_spec(fields, prompt_text)
        fields.pop("assetFolderId", None)
        fields.pop("assetLabel", None)
        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",
        }
        result = self.request_json(
            "POST",
            "/api/external/v1/editor/character-animations/generations",
            body,
            timeout=GENERATION_REQUEST_TIMEOUT_SECONDS,
        )
        if isinstance(session, dict) and isinstance(result, dict) and not result.get("asset"):
            frames = result.get("frames")
            first_frame = frames[0] if isinstance(frames, list) and frames else None
            folder_id = normalize_optional_text(session.get("assetFolderId"))
            if isinstance(first_frame, dict) and folder_id:
                asset = self.create_asset(
                    folder_id,
                    asset_label,
                    first_frame["imageSrc"],
                    int(first_frame["width"]),
                    int(first_frame["height"]),
                    prompt=result.get("prompt"),
                    model=result.get("model"),
                    provider="ark",
                    taskId=result.get("taskId"),
                    assetKind="character-animation",
                    generationInputs={
                        "frames": frames,
                        "previewVideoPath": result.get("previewVideoPath"),
                    },
                )
                result["asset"] = unwrap_envelope(asset).get("asset")
        return result

    def generate_video(self, prompt: str, **fields: Any) -> Any:
        fields.pop("mode", None)
        self._apply_canvas_session_fields(fields, prompt, 1280, 720)
        prompt = self._apply_art_spec(fields, prompt)
        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,
            timeout=GENERATION_REQUEST_TIMEOUT_SECONDS,
        )

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

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


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}
    client = GenarrativeExternalClient(api_key="test")
    session = {"projectId": "proj-demo", "assetFolderId": "editor-asset-folder-demo"}
    fields = client.canvas_generation_fields(session, "英雄角色", width=512, height=768)
    assert fields["projectId"] == "proj-demo"
    assert fields["assetFolderId"] == "editor-asset-folder-demo"
    assert fields["assetLabel"] == "英雄角色"
    assert fields["canvasCompletion"]["placeholder"]["height"] == 768
    assert "美术规范" in art_spec_prompt("生成角色", {"style": "水彩"})
    calls: list[dict[str, Any]] = []

    def fake_request_json(
        method: str,
        path: str,
        body: dict[str, Any] | None = None,
        query: dict[str, Any] | None = None,
        auth: bool = True,
        timeout: int = DEFAULT_REQUEST_TIMEOUT_SECONDS,
    ) -> Any:
        calls.append({"method": method, "path": path, "body": body, "timeout": timeout})
        if path == "/api/external/v1/editor/assets":
            return {"asset": {"assetId": "editor-asset-demo"}}
        return {
            "taskId": "task-demo",
            "model": "seedance2.0-fast",
            "prompt": "角色呼吸",
            "previewVideoPath": "/generated/preview.mp4",
            "frames": [{"frameIndex": 1, "imageSrc": "/generated/frame01.png", "width": 512, "height": 768}],
        }

    client.request_json = fake_request_json  # type: ignore[method-assign]
    result = client.animate_character(
        "/generated/source.png",
        512,
        768,
        "角色呼吸",
        "layer-hero",
        canvasSession=session,
        canvasTitle="角色呼吸动画",
    )
    assert calls[0]["timeout"] == GENERATION_REQUEST_TIMEOUT_SECONDS
    assert calls[0]["body"]["projectId"] == "proj-demo"
    assert calls[0]["body"]["canvasCompletion"]["title"] == "角色呼吸动画"
    assert calls[1]["path"] == "/api/external/v1/editor/assets"
    assert result["asset"]["assetId"] == "editor-asset-demo"
    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())