"""PP-DocLayoutV3 layout-detection server — ONNX Runtime on an NVIDIA L4.

Design notes for the L4 (Ada, 24 GB, 121 TFLOPS fp16, PCIe gen4 x16):
  * TensorRT EP with fp16 is the fast path; Ada has no fp64 units worth using and
    the graph is fp32-in/fp32-out, so fp16 kernels are close to free accuracy-wise.
    CUDA EP is the fallback if TensorRT isn't installed.
  * Engines are cached on disk — a cold TRT build for this graph takes minutes,
    and without the cache you pay it on every process start.
  * Requests are batched dynamically, then padded up to a fixed bucket so TRT only
    ever sees a handful of input shapes (each new shape triggers a rebuild).
  * Device buffers are allocated once per bucket and reused via IOBinding, so the
    steady state does no device allocation.
  * `out_masks` is 48 MB per image in fp32. At bucket 8 that's ~390 MB moved back
    over PCIe per batch (~30-40 ms). If you don't need polygons, serve the model
    exported with --no-masks; the server detects it automatically.

Run:
    python serve_pp_doclayout_v3.py --onnx pp_doclayoutv3.onnx --provider tensorrt
    curl -F "file=@page.jpg" http://localhost:8000/v1/layout
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import binascii
import contextlib
import logging
import os
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Sequence

import numpy as np
import onnxruntime as ort
import uvicorn
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
from fastapi.responses import JSONResponse, PlainTextResponse
from pydantic import BaseModel, Field

from pp_doclayout_v3_onnx import INPUT_SIZE, load_image_rgb, postprocess, preprocess

logger = logging.getLogger("pp-doclayout")

# Bucket sizes the engine is built for. Keep this list short: every distinct
# batch dimension makes TensorRT build (and cache) another engine.
DEFAULT_BUCKETS = (1, 2, 4, 8)


# --------------------------------------------------------------------------- #
# Engine
# --------------------------------------------------------------------------- #
class LayoutEngine:
    """One ORT session plus reusable pinned/device buffers, keyed by batch bucket."""

    def __init__(
        self,
        onnx_path: str,
        *,
        provider: str = "tensorrt",
        device_id: int = 0,
        buckets: Sequence[int] = DEFAULT_BUCKETS,
        trt_cache: str = "./trt_cache",
        trt_fp16: bool = True,
        gpu_mem_limit_gb: float = 20.0,
        intra_op_threads: int = 0,
    ) -> None:
        self.buckets = tuple(sorted(buckets))
        self.max_batch = self.buckets[-1]
        self.device_id = device_id

        so = ort.SessionOptions()
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        if intra_op_threads:
            so.intra_op_num_threads = intra_op_threads
        # The GPU does the work; ORT's CPU-side spinning just burns cores that the
        # preprocessing pool needs.
        so.add_session_config_entry("session.intra_op.allow_spinning", "0")

        providers = self._build_providers(provider, trt_cache, trt_fp16, gpu_mem_limit_gb)
        t0 = time.perf_counter()
        self.session = ort.InferenceSession(onnx_path, sess_options=so, providers=providers)
        self.load_seconds = time.perf_counter() - t0

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        self.has_masks = "out_masks" in self.output_names
        self.active_providers = self.session.get_providers()
        self.on_gpu = any(p in self.active_providers for p in
                          ("TensorrtExecutionProvider", "CUDAExecutionProvider"))

        logger.info(
            "session ready in %.1fs | providers=%s | masks=%s",
            self.load_seconds, self.active_providers, self.has_masks,
        )

        self._device_inputs: dict[int, ort.OrtValue] = {}
        self._bindings: dict[int, Any] = {}
        if self.on_gpu:
            self._alloc_buffers()

    def _build_providers(
        self, provider: str, trt_cache: str, trt_fp16: bool, gpu_mem_limit_gb: float
    ) -> list[Any]:
        available = ort.get_available_providers()
        providers: list[Any] = []

        if provider == "tensorrt":
            if "TensorrtExecutionProvider" not in available:
                logger.warning("TensorRT EP unavailable (have: %s) — falling back to CUDA", available)
            else:
                os.makedirs(trt_cache, exist_ok=True)
                shapes = self._trt_profile_shapes()
                providers.append((
                    "TensorrtExecutionProvider",
                    {
                        "device_id": self.device_id,
                        "trt_fp16_enable": trt_fp16,
                        "trt_engine_cache_enable": True,
                        "trt_engine_cache_path": trt_cache,
                        "trt_timing_cache_enable": True,
                        "trt_timing_cache_path": trt_cache,
                        # 24 GB card; leave headroom for weights, buffers and CUDA ctx.
                        "trt_max_workspace_size": int(6 * 1024**3),
                        "trt_builder_optimization_level": 3,
                        # Deformable attention lands in small subgraphs; letting them
                        # fall back to CUDA is usually faster than forcing them in.
                        "trt_min_subgraph_size": 5,
                        **shapes,
                    },
                ))

        if provider in ("tensorrt", "cuda"):
            providers.append((
                "CUDAExecutionProvider",
                {
                    "device_id": self.device_id,
                    "gpu_mem_limit": int(gpu_mem_limit_gb * 1024**3),
                    "arena_extend_strategy": "kSameAsRequested",
                    # EXHAUSTIVE costs a long first run; HEURISTIC is within a few
                    # percent for a conv stack this shallow.
                    "cudnn_conv_algo_search": "HEURISTIC",
                    "do_copy_in_default_stream": True,
                    "cudnn_conv_use_max_workspace": "1",
                },
            ))

        providers.append("CPUExecutionProvider")
        return providers

    def _trt_profile_shapes(self) -> dict[str, str]:
        lo, hi = self.buckets[0], self.buckets[-1]
        opt = self.buckets[len(self.buckets) // 2]
        fmt = f"{{n}}x3x{INPUT_SIZE}x{INPUT_SIZE}"
        return {
            "trt_profile_min_shapes": f"pixel_values:{fmt.format(n=lo)}",
            "trt_profile_opt_shapes": f"pixel_values:{fmt.format(n=opt)}",
            "trt_profile_max_shapes": f"pixel_values:{fmt.format(n=hi)}",
        }

    def _alloc_buffers(self) -> None:
        """One device input buffer + one binding per bucket, allocated up front."""
        for n in self.buckets:
            host = np.zeros((n, 3, INPUT_SIZE, INPUT_SIZE), dtype=np.float32)
            dev = ort.OrtValue.ortvalue_from_numpy(host, "cuda", self.device_id)
            binding = self.session.io_binding()
            binding.bind_ortvalue_input(self.input_name, dev)
            for name in self.output_names:
                binding.bind_output(name, "cuda", self.device_id)
            self._device_inputs[n] = dev
            self._bindings[n] = binding

    def bucket_for(self, n: int) -> int:
        for b in self.buckets:
            if n <= b:
                return b
        raise ValueError(f"batch {n} exceeds max bucket {self.max_batch}")

    def infer(self, batch: np.ndarray) -> dict[str, np.ndarray]:
        """batch: (n, 3, 800, 800) float32. Returns outputs trimmed back to n."""
        n = batch.shape[0]
        bucket = self.bucket_for(n)

        if not self.on_gpu:  # CPU fallback: no padding, no binding
            outputs = self.session.run(None, {self.input_name: np.ascontiguousarray(batch)})
            return dict(zip(self.output_names, outputs))

        padded = batch
        if n < bucket:  # pad with zeros; the extra rows are discarded below
            padded = np.zeros((bucket, *batch.shape[1:]), dtype=np.float32)
            padded[:n] = batch

        self._device_inputs[bucket].update_inplace(np.ascontiguousarray(padded))
        binding = self._bindings[bucket]
        self.session.run_with_iobinding(binding)
        outputs = [v.numpy() for v in binding.get_outputs()]
        return {name: out[:n] for name, out in zip(self.output_names, outputs)}

    def warmup(self, rounds: int = 2) -> None:
        """Touch every bucket so TRT builds/loads all engines before traffic lands."""
        for n in self.buckets:
            dummy = np.zeros((n, 3, INPUT_SIZE, INPUT_SIZE), dtype=np.float32)
            for _ in range(rounds):
                t0 = time.perf_counter()
                self.infer(dummy)
                logger.info("warmup batch=%d  %.0f ms", n, (time.perf_counter() - t0) * 1e3)


# --------------------------------------------------------------------------- #
# Metrics
# --------------------------------------------------------------------------- #
class Metrics:
    def __init__(self, window: int = 512) -> None:
        self.requests = 0
        self.images = 0
        self.errors = 0
        self.batches = 0
        self.batch_sizes: deque[int] = deque(maxlen=window)
        self.latency_ms: deque[float] = deque(maxlen=window)
        self.infer_ms: deque[float] = deque(maxlen=window)
        self.queue_depth = 0

    @staticmethod
    def _pct(values: Sequence[float], p: float) -> float:
        return float(np.percentile(values, p)) if values else 0.0

    def render(self) -> str:
        lat, inf = list(self.latency_ms), list(self.infer_ms)
        lines = [
            "# HELP layout_requests_total Requests served",
            "# TYPE layout_requests_total counter",
            f"layout_requests_total {self.requests}",
            "# HELP layout_images_total Images processed",
            "# TYPE layout_images_total counter",
            f"layout_images_total {self.images}",
            "# HELP layout_errors_total Failed requests",
            "# TYPE layout_errors_total counter",
            f"layout_errors_total {self.errors}",
            "# HELP layout_batches_total Inference batches executed",
            "# TYPE layout_batches_total counter",
            f"layout_batches_total {self.batches}",
            "# HELP layout_queue_depth Images waiting in the batch queue",
            "# TYPE layout_queue_depth gauge",
            f"layout_queue_depth {self.queue_depth}",
            "# HELP layout_batch_size_avg Mean images per inference batch",
            "# TYPE layout_batch_size_avg gauge",
            f"layout_batch_size_avg {np.mean(self.batch_sizes) if self.batch_sizes else 0:.2f}",
        ]
        for name, vals in (("latency", lat), ("infer", inf)):
            for p in (50, 95, 99):
                lines += [
                    f"# TYPE layout_{name}_ms_p{p} gauge",
                    f"layout_{name}_ms_p{p} {self._pct(vals, p):.2f}",
                ]
        return "\n".join(lines) + "\n"


# --------------------------------------------------------------------------- #
# Dynamic batching
# --------------------------------------------------------------------------- #
@dataclass
class WorkItem:
    """A single image awaiting inference, tied back to its parent request."""

    pixels: np.ndarray                 # (3, 800, 800) float32
    target_size: tuple[int, int]       # (h, w) of the original image
    threshold: float
    future: asyncio.Future = field(repr=False)


class BatchScheduler:
    """Collects work items into batches: fires at max_batch or after max_wait_ms."""

    def __init__(
        self,
        engine: LayoutEngine,
        metrics: Metrics,
        *,
        max_wait_ms: float = 8.0,
        post_pool: Any = None,
    ) -> None:
        self.engine = engine
        self.metrics = metrics
        self.max_wait = max_wait_ms / 1000.0
        self.queue: asyncio.Queue[WorkItem] = asyncio.Queue()
        self.post_pool = post_pool
        self._task: asyncio.Task | None = None

    @property
    def running(self) -> bool:
        return self._task is not None and not self._task.done()

    def start(self) -> None:
        self._task = asyncio.create_task(self._loop())

    async def stop(self) -> None:
        if self._task:
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task

    async def submit(self, item: WorkItem) -> Any:
        await self.queue.put(item)
        self.metrics.queue_depth = self.queue.qsize()
        return await item.future

    async def _collect(self) -> list[WorkItem]:
        first = await self.queue.get()
        items = [first]
        deadline = time.perf_counter() + self.max_wait
        while len(items) < self.engine.max_batch:
            remaining = deadline - time.perf_counter()
            if remaining <= 0:
                break
            try:
                items.append(await asyncio.wait_for(self.queue.get(), timeout=remaining))
            except asyncio.TimeoutError:
                break
        return items

    async def _loop(self) -> None:
        loop = asyncio.get_running_loop()
        while True:
            items: list[WorkItem] = []
            try:
                items = await self._collect()
                self.metrics.queue_depth = self.queue.qsize()
                batch = np.stack([it.pixels for it in items])

                t0 = time.perf_counter()
                # ORT releases the GIL during Run, so the event loop keeps serving.
                outputs = await loop.run_in_executor(None, self.engine.infer, batch)
                infer_ms = (time.perf_counter() - t0) * 1e3

                self.metrics.batches += 1
                self.metrics.batch_sizes.append(len(items))
                self.metrics.infer_ms.append(infer_ms)

                await self._scatter(loop, items, outputs, infer_ms)
            except asyncio.CancelledError:
                raise
            except Exception as exc:  # never let the loop die
                logger.exception("batch failed: %s", exc)
                for it in items:
                    if not it.future.done():
                        it.future.set_exception(exc)

    async def _scatter(self, loop, items, outputs, infer_ms: float) -> None:
        """Post-process each image on the thread pool (cv2 contours release the GIL)."""
        masks = outputs.get("out_masks")

        async def finish(i: int, item: WorkItem) -> None:
            try:
                dets = await loop.run_in_executor(
                    self.post_pool,
                    postprocess,
                    outputs["logits"][i : i + 1],
                    outputs["pred_boxes"][i : i + 1],
                    outputs["order_logits"][i : i + 1],
                    None if masks is None else masks[i : i + 1],
                    [item.target_size],
                    item.threshold,
                )
                if not item.future.done():
                    item.future.set_result((dets[0], infer_ms))
            except Exception as exc:
                if not item.future.done():
                    item.future.set_exception(exc)

        await asyncio.gather(*(finish(i, it) for i, it in enumerate(items)))


# --------------------------------------------------------------------------- #
# API
# --------------------------------------------------------------------------- #
class Base64Request(BaseModel):
    images: list[str] = Field(..., description="base64-encoded image bytes")
    threshold: float = Field(0.5, ge=0.0, le=1.0)


def serialise(detections: list[dict], include_polygons: bool) -> list[dict]:
    out = []
    for d in detections:
        item = {
            "order": d["order"],
            "label": d["label"],
            "label_id": d["label_id"],
            "score": round(d["score"], 4),
            "box": d["box"],
        }
        if include_polygons:
            item["polygon"] = np.asarray(d["polygon"]).round(1).tolist()
        out.append(item)
    return out


def build_app(engine: LayoutEngine, scheduler: BatchScheduler, metrics: Metrics,
              pre_pool: Any, post_pool: Any = None) -> FastAPI:
    @contextlib.asynccontextmanager
    async def lifespan(_: FastAPI):
        scheduler.start()
        logger.info("scheduler started; accepting traffic")
        yield
        await scheduler.stop()
        for pool in (pre_pool, post_pool):
            if pool is not None:
                pool.shutdown(wait=False)

    app = FastAPI(title="PP-DocLayoutV3", version="1.0", lifespan=lifespan)

    async def run_images(raw: list[bytes], threshold: float) -> list[list[dict]]:
        loop = asyncio.get_running_loop()

        def decode(blob: bytes):
            arr = np.frombuffer(blob, dtype=np.uint8)
            rgb = load_image_rgb(_imdecode(arr))
            pixels, sizes = preprocess([rgb])
            return pixels[0], sizes[0]

        try:
            prepared = await asyncio.gather(
                *(loop.run_in_executor(pre_pool, decode, blob) for blob in raw)
            )
        except Exception as exc:
            raise HTTPException(status_code=400, detail=f"could not decode image: {exc}") from exc

        results = await asyncio.gather(*(
            scheduler.submit(WorkItem(pixels=p, target_size=s, threshold=threshold,
                                      future=loop.create_future()))
            for p, s in prepared
        ))
        return [r[0] for r in results]

    @app.post("/v1/layout")
    async def layout(
        file: list[UploadFile] = File(...),
        threshold: float = Query(0.5, ge=0.0, le=1.0),
        polygons: bool = Query(True),
    ) -> JSONResponse:
        t0 = time.perf_counter()
        blobs = [await f.read() for f in file]
        if len(blobs) > 64:
            raise HTTPException(413, "at most 64 images per request")
        try:
            detections = await run_images(blobs, threshold)
        except HTTPException:
            metrics.errors += 1
            raise
        except Exception as exc:
            metrics.errors += 1
            logger.exception("inference failed")
            raise HTTPException(500, str(exc)) from exc

        latency_ms = (time.perf_counter() - t0) * 1e3
        metrics.requests += 1
        metrics.images += len(blobs)
        metrics.latency_ms.append(latency_ms)

        want_polygons = polygons and engine.has_masks
        return JSONResponse({
            "latency_ms": round(latency_ms, 2),
            "results": [
                {"index": i, "elements": serialise(d, want_polygons), "count": len(d)}
                for i, d in enumerate(detections)
            ],
        })

    @app.post("/v1/layout/base64")
    async def layout_base64(body: Base64Request) -> JSONResponse:
        try:
            blobs = [base64.b64decode(s, validate=True) for s in body.images]
        except (binascii.Error, ValueError) as exc:
            metrics.errors += 1
            raise HTTPException(400, f"invalid base64: {exc}") from exc
        t0 = time.perf_counter()
        detections = await run_images(blobs, body.threshold)
        metrics.requests += 1
        metrics.images += len(blobs)
        latency_ms = (time.perf_counter() - t0) * 1e3
        metrics.latency_ms.append(latency_ms)
        return JSONResponse({
            "latency_ms": round(latency_ms, 2),
            "results": [
                {"index": i, "elements": serialise(d, engine.has_masks), "count": len(d)}
                for i, d in enumerate(detections)
            ],
        })

    @app.get("/healthz")
    async def healthz() -> dict:
        return {"status": "ok"}

    @app.get("/readyz")
    async def readyz() -> dict:
        if not scheduler.running:
            raise HTTPException(503, "scheduler not running")
        return {"status": "ready"}

    @app.get("/info")
    async def info() -> dict:
        return {
            "providers": engine.active_providers,
            "on_gpu": engine.on_gpu,
            "masks": engine.has_masks,
            "buckets": list(engine.buckets),
            "input_size": INPUT_SIZE,
            "load_seconds": round(engine.load_seconds, 2),
        }

    @app.get("/metrics")
    async def prometheus() -> PlainTextResponse:
        return PlainTextResponse(metrics.render())

    return app


def _imdecode(arr: np.ndarray) -> np.ndarray:
    import cv2

    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("unsupported or corrupt image data")
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


def main() -> int:
    from concurrent.futures import ThreadPoolExecutor

    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--onnx", required=True)
    p.add_argument("--provider", default="tensorrt", choices=["tensorrt", "cuda", "cpu"])
    p.add_argument("--device-id", type=int, default=0)
    p.add_argument("--buckets", default="1,2,4,8", help="batch buckets the engine is built for")
    p.add_argument("--max-wait-ms", type=float, default=8.0, help="how long to wait to fill a batch")
    p.add_argument("--trt-cache", default="./trt_cache")
    p.add_argument("--no-fp16", dest="fp16", action="store_false")
    p.add_argument("--gpu-mem-gb", type=float, default=20.0)
    p.add_argument("--pre-threads", type=int, default=4, help="decode/resize workers")
    p.add_argument("--post-threads", type=int, default=4, help="polygon extraction workers")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument("--log-level", default="info")
    args = p.parse_args()

    logging.basicConfig(
        level=args.log_level.upper(),
        format="%(asctime)s %(levelname)s %(name)s | %(message)s",
    )
    # OpenCV spawning threads inside our own pool oversubscribes the CPU.
    import cv2

    cv2.setNumThreads(1)

    engine = LayoutEngine(
        args.onnx,
        provider=args.provider,
        device_id=args.device_id,
        buckets=[int(b) for b in args.buckets.split(",")],
        trt_cache=args.trt_cache,
        trt_fp16=args.fp16,
        gpu_mem_limit_gb=args.gpu_mem_gb,
    )
    engine.warmup()

    metrics = Metrics()
    pre_pool = ThreadPoolExecutor(args.pre_threads, thread_name_prefix="pre")
    post_pool = ThreadPoolExecutor(args.post_threads, thread_name_prefix="post")
    scheduler = BatchScheduler(engine, metrics, max_wait_ms=args.max_wait_ms, post_pool=post_pool)

    app = build_app(engine, scheduler, metrics, pre_pool, post_pool)
    logger.info("serving on %s:%d", args.host, args.port)

    # Single process on purpose: one CUDA context, one set of TRT engines.
    uvicorn.run(app, host=args.host, port=args.port, log_level=args.log_level, workers=1)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())