"""JarvisChat routers — Image generation proxy endpoint."""
import base64
import logging

from fastapi import APIRouter, HTTPException
from fastapi.responses import Response

from cluster import CLUSTER_NODES, request_image_generate

log = logging.getLogger("caic")
router = APIRouter()


def _find_image_node() -> str | None:
    for name, node in CLUSTER_NODES.items():
        if node.get("status") == "active" and "image_gen" in node.get("capabilities", []):
            return name
    return None


@router.post("/api/image/generate")
async def generate_image(request_body: dict):
    prompt = (request_body.get("prompt") or "").strip()
    if not prompt:
        raise HTTPException(status_code=400, detail="Prompt is required")

    negative_prompt = request_body.get("negative_prompt", "")
    width = min(max(request_body.get("width", 1024), 256), 2048)
    height = min(max(request_body.get("height", 1024), 256), 2048)
    steps = min(max(request_body.get("steps", 20), 1), 50)
    seed = request_body.get("seed", -1)
    model = request_body.get("model", "")

    node_name = _find_image_node()
    if not node_name:
        raise HTTPException(status_code=503, detail="No image generation service available")

    log.info("image generate via %s: %s", node_name, prompt[:60])

    image_b64 = await request_image_generate(
        node_name=node_name,
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        steps=steps,
        seed=seed,
        model=model,
    )

    if image_b64 is None:
        raise HTTPException(status_code=504, detail="Image generation timed out or failed")

    image_bytes = base64.b64decode(image_b64)
    return Response(content=image_bytes, media_type="image/png")


@router.get("/api/image/status")
async def image_status():
    nodes = []
    for name, node in CLUSTER_NODES.items():
        caps = node.get("capabilities", [])
        if "image_gen" in caps:
            nodes.append({
                "name": name,
                "status": node.get("status"),
                "load": node.get("load"),
                "last_seen": node.get("last_seen"),
            })
    return {"available": len(nodes) > 0, "nodes": nodes}