import io

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from pydantic import BaseModel

from server.inpainting import Inpainter, composite_inpaint
from server.segmentation import Segmenter, resize_for_upload
from server.session_store import SessionStore


class Point(BaseModel):
    x: int
    y: int
    positive: bool


class SegmentPointRequest(BaseModel):
    session_id: str
    points: list[Point]


class InpaintRequest(BaseModel):
    session_id: str


def _png_response(image: Image.Image):
    from fastapi.responses import Response

    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return Response(content=buffer.getvalue(), media_type="image/png")


def create_fastapi_app(
    segmenter: Segmenter,
    inpainter: Inpainter,
    session_store: SessionStore,
) -> FastAPI:
    app = FastAPI()
    # Lets the Flutter web build (served from a browser origin) call this
    # API directly - browsers block cross-origin requests by default, and
    # this server has no cookies/session-cookie auth to protect.
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @app.get("/health")
    def health():
        return {"status": "ok"}

    @app.post("/segment/init")
    async def segment_init(image: UploadFile = File(...)):
        contents = await image.read()
        pil_image = Image.open(io.BytesIO(contents)).convert("RGB")
        pil_image = resize_for_upload(pil_image)
        segmenter.encode(pil_image)
        session_id = session_store.create({"image": pil_image})
        return {"session_id": session_id, "width": pil_image.width, "height": pil_image.height}

    @app.post("/segment/point")
    def segment_point(request: SegmentPointRequest):
        session = session_store.get(request.session_id)
        if session is None:
            raise HTTPException(status_code=404, detail="session not found")
        points = [(p.x, p.y, p.positive) for p in request.points]
        mask = segmenter.decode(session["image"], points)
        return _png_response(mask)

    @app.post("/inpaint")
    async def inpaint(
        image: UploadFile = File(...),
        mask: UploadFile = File(...),
    ):
        image_bytes = await image.read()
        mask_bytes = await mask.read()
        pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        pil_mask = Image.open(io.BytesIO(mask_bytes)).convert("L")
        if pil_mask.size != pil_image.size:
            pil_mask = pil_mask.resize(pil_image.size)
        if inpainter.does_own_compositing:
            result = inpainter.inpaint(pil_image, pil_mask)
        else:
            result = composite_inpaint(pil_image, pil_mask, inpainter)
        return _png_response(result)

    return app


# ZeroGPU requires @spaces.GPU to be a real decorator on a module-level
# function - assigning `obj.method = spaces.GPU(obj.method)` at runtime
# (the previous approach) is not detected during Space startup ("No
# @spaces.GPU function detected during startup"). `spaces` isn't installed
# in the local dev/test environment (only on the real HF Space runtime), so
# this whole block is guarded - the test suite only ever uses
# create_fastapi_app() with fakes and never imports this branch's names.
try:
    import spaces

    @spaces.GPU
    def _gpu_encode(image):
        from server.model_loader import get_segmenter

        return get_segmenter().encode(image)

    @spaces.GPU
    def _gpu_decode(image, points):
        from server.model_loader import get_segmenter

        return get_segmenter().decode(image, points)

    @spaces.GPU
    def _gpu_inpaint(image, mask):
        from server.model_loader import get_inpainter

        return get_inpainter().inpaint(image, mask)

except ImportError:
    spaces = None


class _GpuSegmenter(Segmenter):
    """Delegates to the module-level @spaces.GPU-decorated functions above."""

    def encode(self, image):
        return _gpu_encode(image)

    def decode(self, image, points):
        return _gpu_decode(image, points)


class _GpuInpainter(Inpainter):
    """Delegates to the module-level @spaces.GPU-decorated function above."""

    def inpaint(self, image, mask):
        return _gpu_inpaint(image, mask)


def build_production_app():
    """Entrypoint used by the Hugging Face Space runtime (`python server/app.py`)."""
    import gradio as gr

    from server.session_store import SessionStore

    session_store = SessionStore(ttl_seconds=600)
    fastapi_app = create_fastapi_app(_GpuSegmenter(), _GpuInpainter(), session_store)

    with gr.Blocks() as demo:
        gr.Markdown(
            "# Object Removal API\n"
            "This Space exposes `/health`, `/segment/init`, `/segment/point`, "
            "`/inpaint` for the Object Removal mobile app. No UI is served here."
        )
        # ZeroGPU's startup handshake (spaces/zero/gradio.py's one_launch)
        # monkey-patches gr.Blocks.launch itself: the very first call to
        # .launch() on any Blocks instance runs a startup() function that
        # reports the @spaces.GPU-decorated functions to HF's ZeroGPU infra.
        # gr.mount_gradio_app() + a manually-run uvicorn server (the previous
        # approach here) never calls .launch() at all, so that handshake
        # never fires and the infra times out waiting for it ("No
        # @spaces.GPU function detected during startup") - the message is
        # misleading, it actually means "no startup report was ever sent".
        # _gpu_inpaint is wired directly as a Gradio fn= here (Image+Image ->
        # Image maps cleanly onto gr.Image components) so the decorator scan
        # that runs during .launch() finds it; the real image data still
        # flows through the /segment/init, /segment/point, /inpaint FastAPI
        # routes above, which call the same underlying _gpu_encode/
        # _gpu_decode/_gpu_inpaint functions directly.
        with gr.Row(visible=False):
            gr.Interface(
                fn=_gpu_inpaint,
                inputs=[gr.Image(type="pil"), gr.Image(type="pil", image_mode="L")],
                outputs=gr.Image(type="pil"),
            )

    # demo.launch() (not gr.mount_gradio_app) is required so ZeroGPU's
    # startup handshake fires. prevent_thread_lock=True lets it start its
    # own uvicorn server in a background thread and return immediately, so
    # our custom FastAPI routes can be merged onto the same app afterwards.
    # ssr_mode=False is required too: HF Spaces sets GRADIO_SSR_MODE=true by
    # default, which installs a middleware that proxies every non-internal
    # path (including our /health etc. and even Gradio's own /docs) to a
    # Node SSR renderer before Starlette routing ever runs - every path
    # silently returned the same SPA shell HTML until this was disabled.
    served_app, _local_url, _share_url = demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        prevent_thread_lock=True,
        show_error=True,
        quiet=True,
        ssr_mode=False,
    )
    served_app.include_router(fastapi_app.router)
    return demo


if __name__ == "__main__":
    demo = build_production_app()
    demo.block_thread()