render_config.py

"""渲染配置：严格校验所有来自 UI、API 和 MCP 的 JSON 参数。"""

from __future__ import annotations

import json
import os
import re
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


MAX_CONFIG_BYTES = 64 * 1024
_DRIVE = re.compile(r"^[A-Za-z]:")


def positive_env_int(name: str, default: int) -> int:
    """读取服务端资源上限；错误配置直接报错，避免意外取消限制。"""
    value = os.environ.get(name, str(default))
    try:
        parsed = int(value)
    except (ValueError, TypeError) as exc:
        raise ValueError(f"环境变量 {name} 必须为正整数。") from exc
    if parsed <= 0:
        raise ValueError(f"环境变量 {name} 必须为正整数。")
    return parsed


def safe_relative_path(value: str) -> str:
    """保留路径原样；禁止跨平台歧义、越界路径和 Windows ADS。"""
    if not isinstance(value, str) or not value:
        raise ValueError("文件路径必须是非空相对路径。")
    if value.startswith("/") or _DRIVE.match(value) or "\\" in value:
        raise ValueError("文件路径必须使用相对路径和 / 分隔符。")
    if any(ord(char) < 32 or ord(char) == 127 for char in value):
        raise ValueError("文件路径不能包含 NUL 或控制字符。")
    parts = value.split("/")
    if any(part in {"", ".", ".."} for part in parts):
        raise ValueError("文件路径不能包含空层级、. 或 ..。")
    if any(":" in part or part.endswith((".", " ")) for part in parts):
        raise ValueError("文件路径包含不安全的跨平台文件名。")
    return value


class RenderConfig(BaseModel):
    """可直接传递给渲染任务的、版本化的严格配置。"""

    model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False)

    schema_version: Literal[1] = 1
    blend_file: str | None = None
    mode: Literal["image", "frames", "video"] = "image"
    scene: str | None = None
    camera: str | None = None
    view_layer: str | None = None
    width: Annotated[int, Field(ge=16, le=8192)] = 1280
    height: Annotated[int, Field(ge=16, le=8192)] = 720
    resolution_percentage: Annotated[int, Field(ge=1, le=100)] = 100
    samples: Annotated[int, Field(ge=1, le=4096)] = 64
    denoise: bool = True
    adaptive_threshold: Annotated[float, Field(ge=0, le=1)] = 0.02
    transparent: bool = False
    frame: Annotated[int, Field(ge=1, le=1_000_000)] = 1
    frame_start: Annotated[int, Field(ge=1, le=1_000_000)] = 1
    frame_end: Annotated[int, Field(ge=1, le=1_000_000)] = 120
    frame_step: Annotated[int, Field(ge=1, le=1_000_000)] = 1
    fps: Annotated[int, Field(ge=1, le=120)] = 24
    format: Literal["PNG", "OPEN_EXR"] = "PNG"
    color_depth: Literal[8, 16, 32] = 8
    engine: Literal["EEVEE", "CYCLES"] = "EEVEE"
    device: Literal["AUTO", "CUDA", "OPTIX", "CPU"] = "AUTO"
    gpu_size: Literal["large", "xlarge"] = "large"
    graphics_backend: Literal["opengl", "vulkan"] = "opengl"
    seed: Annotated[int, Field(ge=0, le=2_147_483_647)] = 0
    motion_blur: bool = False
    strict_assets: bool = True
    batch_size: Annotated[int, Field(ge=1, le=64)] = 1
    gpu_seconds: Annotated[int, Field(ge=10, le=600)] = 120
    cpu_timeout: Annotated[int, Field(ge=10, le=7200)] = 600
    video_crf: Annotated[int, Field(ge=0, le=40)] = 18
    video_preset: Literal["medium", "fast", "slow"] = "medium"

    @field_validator("schema_version", "color_depth", mode="before")
    @classmethod
    def require_integer_literal(cls, value: object) -> object:
        # Python 的 True == 1；Literal 仍需显式防止 bool/float 伪装成整数。
        if type(value) is not int:
            raise ValueError("此字段必须使用 JSON 整数。")
        return value

    @field_validator("blend_file")
    @classmethod
    def validate_blend_file(cls, value: str | None) -> str | None:
        if value is not None:
            safe_relative_path(value)
            if not value.casefold().endswith(".blend"):
                raise ValueError("blend_file 必须指定 .blend 文件。")
        return value

    @field_validator("scene", "camera", "view_layer")
    @classmethod
    def validate_datablock_name(cls, value: str | None) -> str | None:
        if value is not None and (
            not value.strip() or len(value) > 1024 or "\x00" in value
        ):
            raise ValueError("场景、相机和视图层名称必须为有效的非空名称。")
        return value

    @model_validator(mode="after")
    def validate_render_constraints(self) -> "RenderConfig":
        if self.engine == "EEVEE" and self.device in {"CUDA", "OPTIX"}:
            raise ValueError("CUDA / OPTIX 仅用于 Cycles；Eevee 请使用 CPU 或 AUTO。")
        if self.frame_end < self.frame_start:
            raise ValueError("frame_end 不能小于 frame_start。")
        if self.format == "PNG" and self.color_depth not in (8, 16):
            raise ValueError("PNG 的 color_depth 只能为 8 或 16。")
        if self.format == "OPEN_EXR" and self.color_depth not in (16, 32):
            raise ValueError("OPEN_EXR 的 color_depth 只能为 16 或 32。")
        if self.output_width < 1 or self.output_height < 1:
            raise ValueError("缩放后的输出尺寸必须至少为 1 × 1。")
        if self.output_width * self.output_height > positive_env_int(
            "MAX_PIXELS", 16_777_216
        ):
            raise ValueError("输出像素总数超过服务端 MAX_PIXELS 限制。")
        frame_count = 1 if self.mode == "image" else (
            (self.frame_end - self.frame_start) // self.frame_step + 1
        )
        if frame_count > positive_env_int("MAX_FRAMES", 2000):
            raise ValueError("帧数超过服务端 MAX_FRAMES 限制。")
        if self.mode == "video":
            if self.format != "PNG":
                raise ValueError("video 模式只支持 PNG 中间帧。")
            if self.transparent:
                raise ValueError("MP4 视频不支持透明背景，请设置 transparent=false。")
            if self.output_width % 2 or self.output_height % 2:
                raise ValueError("video 模式的最终宽、高必须均为偶数。")
        return self

    @property
    def output_width(self) -> int:
        return max(1, self.width * self.resolution_percentage // 100)

    @property
    def output_height(self) -> int:
        return max(1, self.height * self.resolution_percentage // 100)

    def frame_numbers(self) -> list[int]:
        """逐帧任务使用此列表；隔帧采样由编码端同步调整播放帧率以保持时间尺度。"""
        if self.mode == "image":
            return [self.frame]
        return list(range(self.frame_start, self.frame_end + 1, self.frame_step))


def _reject_duplicate_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
    result: dict[str, object] = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"JSON 包含重复字段：{key}")
        result[key] = value
    return result


def _reject_nonfinite(value: str) -> None:
    raise ValueError(f"JSON 不允许非有限数值：{value}")


def parse_config(text: str) -> RenderConfig:
    """解析 JSON；限制输入大小并拒绝重复字段、NaN 和隐式类型转换。"""
    if not isinstance(text, str):
        raise ValueError("配置必须是 JSON 文本。")
    try:
        byte_count = len(text.encode("utf-8"))
    except UnicodeEncodeError as exc:
        raise ValueError("配置文本包含无效 Unicode 字符。") from exc
    if byte_count > MAX_CONFIG_BYTES:
        raise ValueError("配置文件不能超过 64 KiB。")
    try:
        raw = json.loads(
            text, object_pairs_hook=_reject_duplicate_keys, parse_constant=_reject_nonfinite
        )
    except (json.JSONDecodeError, RecursionError) as exc:
        raise ValueError(f"无效的 JSON 配置：{exc}") from exc
    if not isinstance(raw, dict):
        raise ValueError("配置根节点必须是 JSON 对象。")
    return RenderConfig.model_validate(raw)


def json_schema() -> dict:
    """返回标准 JSON Schema，方便 agent 发现可用参数。"""
    return RenderConfig.model_json_schema()