"""First-party browser UI routes."""

from __future__ import annotations

from base64 import b64encode
from pathlib import Path

from flask import jsonify, redirect, request, send_from_directory

from diffusion_cli.config import (
    ImageGenerationRequest,
    OUTPUT_FORMAT_OPTIONS,
    UI_RESOLUTION_PRESETS,
    buildDefaultGenerationRequest,
    normalizeOutputExtension,
    outputMimeType,
    validateDimensions,
    validateOutputQuality,
)
from diffusion_cli.errors import DiffusionCliError

UI_ROOT = Path(__file__).parent / "static" / "ui"


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 _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 _validateOptionalDimensions(width: int | None, height: int | None) -> None:
    if width is None and height is None:
        return
    effective_width = 8 if width is None else width
    effective_height = 8 if height is None else height
    validateDimensions(effective_width, effective_height)


def _txt2imgRequest(data: dict) -> ImageGenerationRequest:
    prompt = data.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        raise DiffusionCliError("prompt required")

    negative_prompt = _optionalString(data, "negative_prompt")
    if negative_prompt is not None and not negative_prompt.strip():
        negative_prompt = None

    seed = _optionalInt(data, "seed")
    if seed == -1:
        seed = None

    width = _optionalInt(data, "width")
    height = _optionalInt(data, "height")
    _validateOptionalDimensions(width, height)

    batch_size = _optionalInt(data, "batch_size")
    steps = _optionalInt(data, "steps")
    cfg = _optionalFloat(data, "cfg")
    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 batch_size is not None and batch_size < 1:
        raise DiffusionCliError(
            f"Batch size must be at least 1: got {batch_size}"
        )
    if steps is not None and steps < 1:
        raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
    if cfg is not None and cfg < 0:
        raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")

    return ImageGenerationRequest(
        prompt=prompt.strip(),
        negative_prompt=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 _defaultResponse(context) -> dict:
    default_request = buildDefaultGenerationRequest(
        context.generation_service.user_config,
    )
    return {
        "negative_prompt": default_request.negative_prompt,
        "width": default_request.width,
        "height": default_request.height,
        "batch_size": default_request.batch_size,
        "steps": default_request.steps,
        "cfg": default_request.cfg,
        "output_extension": default_request.output_extension,
        "output_quality": default_request.output_quality,
        "output_formats": list(OUTPUT_FORMAT_OPTIONS),
        "seed": default_request.seed,
        "resolution_presets": [
            {
                "width": width,
                "height": height,
                "resolution": resolution,
                "aspect_ratio": aspect_ratio,
                "label": (
                    f"{width}x{height} ({aspect_ratio})"
                ),
            }
            for width, height, resolution, aspect_ratio
            in UI_RESOLUTION_PRESETS
        ],
    }


def registerUiRoutes(app, context) -> None:
    """Register first-party browser UI routes."""

    @app.route("/ui", methods=["GET"])
    def uiRootRedirect():
        return redirect("/ui/")

    @app.route("/ui/", methods=["GET"])
    def uiIndex():
        return send_from_directory(UI_ROOT, "index.html")

    @app.route("/ui/app.js", methods=["GET"])
    def uiApp():
        return send_from_directory(UI_ROOT, "app.js")

    @app.route("/ui/style.css", methods=["GET"])
    def uiStyle():
        return send_from_directory(UI_ROOT, "style.css")

    @app.route("/ui/vendor/preact.module.js", methods=["GET"])
    def uiPreact():
        return send_from_directory(UI_ROOT / "vendor", "preact.module.js")

    @app.route("/ui/vendor/LICENSE.preact", methods=["GET"])
    def uiPreactLicense():
        return send_from_directory(UI_ROOT / "vendor", "LICENSE.preact")

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

    @app.route("/ui/api/defaults", methods=["GET"])
    def uiDefaults():
        return jsonify(_defaultResponse(context))

    @app.route("/ui/api/txt2img", methods=["POST"])
    def uiTxt2Img():
        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")

        return jsonify({
            "images": [
                {
                    "mime_type": outputMimeType(image.format),
                    "data": b64encode(image.data).decode("ascii"),
                    "seed": image.seed,
                }
                for image in images
            ]
        })