BareGit
"""Flask server construction for API profile mode."""

from __future__ import annotations

from dataclasses import dataclass
from types import SimpleNamespace

from flask import Flask, request

from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.generation_service import (
    MODEL_RESIDENCY_VALUES,
    GenerationService,
)
from diffusion_cli.ui import registerUiRoutes

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


@dataclass(frozen=True)
class ServerConfig:
    """Validated settings for HTTP server mode."""

    api_profile: str
    host: str
    port: int
    model_residency: str


def validateServerConfig(
    api_profile: str | None,
    host: str,
    port: int,
    model_residency: str,
) -> ServerConfig:
    """Validate server mode command line settings."""

    if api_profile is None:
        raise DiffusionCliError("--api-profile is required")
    if api_profile not in API_PROFILES:
        raise DiffusionCliError(f"Unknown API profile: {api_profile}")
    if not host:
        raise DiffusionCliError("--host must not be empty")
    if port < 1 or port > 65535:
        raise DiffusionCliError(f"--port must be from 1 through 65535: {port}")
    if model_residency not in MODEL_RESIDENCY_VALUES:
        raise DiffusionCliError(
            "--model-residency must be one of staged, cpu-cache"
        )
    return ServerConfig(
        api_profile=api_profile,
        host=host,
        port=port,
        model_residency=model_residency,
    )


def _summarizeJsonBody(value):
    if isinstance(value, dict):
        result = {}
        for key, item in value.items():
            if key in BINARY_JSON_FIELDS and isinstance(item, str):
                result[key] = f"<{len(item)} bytes>"
            elif key in BINARY_JSON_FIELDS and isinstance(item, list):
                result[key] = f"<{len(item)} items>"
            else:
                result[key] = _summarizeJsonBody(item)
        return result
    if isinstance(value, list):
        return [_summarizeJsonBody(item) for item in value]
    return value


def createApp(server_config: ServerConfig, generation_service) -> Flask:
    """Create a Flask app for the selected API profile."""

    app = Flask(__name__)
    context = SimpleNamespace(
        server_config=server_config,
        generation_service=generation_service,
    )

    @app.before_request
    def logRequest() -> None:
        body = None
        if request.is_json:
            body = _summarizeJsonBody(request.get_json(silent=True))
        app.logger.debug(
            "request method=%s path=%s query=%s json=%r",
            request.method,
            request.path,
            request.query_string.decode("utf-8"),
            body,
        )

    registerUiRoutes(app, context)
    API_PROFILES[server_config.api_profile].register_routes(app, context)
    return app


def serve(server_config: ServerConfig, user_config) -> None:
    """Run the configured Flask development server."""

    generation_service = GenerationService(
        user_config,
        model_residency=server_config.model_residency,
    )
    app = createApp(server_config, generation_service)
    app.run(host=server_config.host, port=server_config.port)