from __future__ import annotations

from typing import Any, Dict

from fastapi import APIRouter

from ..deps import OUTPUT_DIR
from ..schemas import InpaintRequest
from ..services.image_io import b64_to_pil_image, default_half_mask, pil_image_to_png_b64
from ..services.predictors import get_inpaint_predictor

router = APIRouter(tags=["inpaint"])


@router.post("/inpaint")
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
    try:
        model_name = req.model_name or "sdxl_inpaint"
        pil = b64_to_pil_image(req.image_b64).convert("RGB")

        if req.mask_b64:
            mask = b64_to_pil_image(req.mask_b64).convert("L")
        else:
            mask = default_half_mask(pil)

        predictor = get_inpaint_predictor(model_name)
        out = predictor(
            pil,
            mask,
            req.prompt or "",
            req.negative_prompt or "",
            strength=req.strength,
            max_side=req.max_side,
        )

        out_dir = OUTPUT_DIR / "inpaint"
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = out_dir / f"{model_name}_inpaint.png"
        out.save(out_path)

        return {
            "success": True,
            "output_path": str(out_path),
            "output_image_b64": pil_image_to_png_b64(out),
        }
    except Exception as e:
        return {"success": False, "error": str(e)}