BareGit
"""HTTP API profile registry and route implementations."""

from __future__ import annotations

from base64 import b64encode
from dataclasses import dataclass
import json
from typing import Callable

from flask import Response, jsonify, request

from diffusion_cli.config import (
    ImageGenerationRequest,
    normalizeOutputExtension,
    validateDimensions,
    validateOutputQuality,
)
from diffusion_cli.errors import DiffusionCliError

UNSUPPORTED_IMAGE_FIELDS = ("init_images", "mask", "extra_images")


@dataclass(frozen=True)
class ApiProfile:
    """HTTP API contract exposed by server mode."""

    name: str
    description: str
    register_routes: Callable[[object, object], None]


def _optionalString(data: dict, key: str) -> str | None:
    value = data.get(key)
    if value is None:
        return None
    if not isinstance(value, str):
        raise DiffusionCliError(f"{key} must be a string")
    return value


def _optionalInt(data: dict, key: str) -> int | None:
    value = data.get(key)
    if value is None:
        return None
    if not isinstance(value, int) or isinstance(value, bool):
        raise DiffusionCliError(f"{key} must be an integer")
    return value


def _optionalFloat(data: dict, key: str) -> float | None:
    value = data.get(key)
    if value is None:
        return None
    if not isinstance(value, int | float) or isinstance(value, bool):
        raise DiffusionCliError(f"{key} must be a number")
    return float(value)


def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
    for field in UNSUPPORTED_IMAGE_FIELDS:
        if field in data:
            raise DiffusionCliError(f"{field} is not supported")

    prompt = data.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        raise DiffusionCliError("prompt required")

    seed = _optionalInt(data, "seed")
    if seed == -1:
        seed = None
    width = _optionalInt(data, "width")
    height = _optionalInt(data, "height")
    if width is not None and height is not None:
        validateDimensions(width, height)
    elif width is not None and (width <= 0 or width % 8 != 0):
        raise DiffusionCliError(
            f"Width must be a positive multiple of 8: got {width}"
        )
    elif height is not None and (height <= 0 or height % 8 != 0):
        raise DiffusionCliError(
            f"Height must be a positive multiple of 8: got {height}"
        )
    steps = _optionalInt(data, "steps")
    batch_size = _optionalInt(data, "batch_size")
    cfg = _optionalFloat(data, "cfg_scale")
    output_extension = _optionalString(data, "output_extension")
    output_quality = _optionalInt(data, "output_quality")
    if output_extension is not None:
        output_extension = normalizeOutputExtension(output_extension)
    if output_quality is not None:
        output_quality = validateOutputQuality(output_quality)
    if steps is not None and steps < 1:
        raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
    if batch_size is not None and batch_size < 1:
        raise DiffusionCliError(
            f"Batch size must be at least 1: got {batch_size}"
        )
    if cfg is not None and cfg < 0:
        raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")

    return ImageGenerationRequest(
        prompt=prompt,
        negative_prompt=_optionalString(data, "negative_prompt"),
        seed=seed,
        width=width,
        height=height,
        batch_size=batch_size,
        steps=steps,
        cfg=cfg,
        output_extension=output_extension,
        output_quality=output_quality,
    )


def _badRequest(message: str):
    return jsonify({"error": "bad_request", "message": message}), 400


def _generationFailed(message: str):
    return jsonify({"error": "generation_failed", "message": message}), 500


def registerSillyTavernSdcppRoutes(app, context) -> None:
    """Register SillyTavern stable-diffusion.cpp compatibility routes."""

    @app.route("/health", methods=["GET"])
    def health():
        return jsonify({
            "status": "ok",
            "api_profile": context.server_config.api_profile,
        })

    @app.route("/v1/images/generations", methods=["OPTIONS"])
    def imageGenerationsOptions():
        response = Response(status=204)
        response.headers["Allow"] = "OPTIONS, POST"
        return response

    @app.route("/v1/models", methods=["GET"])
    def models():
        return jsonify({
            "data": [
                {
                    "id": "z-image-local",
                    "object": "model",
                    "owned_by": "local",
                }
            ]
        })

    @app.route("/sdapi/v1/txt2img", methods=["POST"])
    def txt2img():
        data = request.get_json(silent=True)
        if not isinstance(data, dict) or not data:
            return _badRequest("JSON object body required")

        try:
            generation_request = _txt2imgRequest(data)
        except DiffusionCliError as exc:
            return _badRequest(str(exc))

        try:
            images = context.generation_service.generateImages(
                generation_request,
            )
        except DiffusionCliError as exc:
            app.logger.exception("Generation failed")
            return _generationFailed(str(exc))
        except Exception:
            app.logger.exception("Generation failed")
            return _generationFailed("generation failed")

        encoded_images = [
            b64encode(image.data).decode("ascii") for image in images
        ]
        info = {
            "prompt": generation_request.prompt,
            "negative_prompt": generation_request.negative_prompt,
            "seed": images[0].seed if images else generation_request.seed,
            "width": generation_request.width,
            "height": generation_request.height,
            "steps": generation_request.steps,
            "cfg_scale": generation_request.cfg,
        }
        return jsonify({
            "images": encoded_images,
            "parameters": data,
            "info": json.dumps(info, separators=(",", ":")),
        })


SILLYTAVERN_SDCPP_PROFILE = ApiProfile(
    name="sillytavern-sdcpp",
    description="SillyTavern stable-diffusion.cpp adapter compatibility.",
    register_routes=registerSillyTavernSdcppRoutes,
)

API_PROFILES = {
    SILLYTAVERN_SDCPP_PROFILE.name: SILLYTAVERN_SDCPP_PROFILE,
}