"""HTTP server: POST /v1/systemone and GET /health, with dynamic micro-batching on one GPU.

    GLIDE_MODEL=/path/to/model uvicorn glide.server:app --host 0.0.0.0 --port 8000
    python -m glide.server --model /path/to/model --port 8000

Run one server process per GPU. The model is loaded once at startup; concurrent requests are queued and scored together in
micro-batches by a single GPU worker thread.

Environment (defaults in brackets):
    GLIDE_MODEL               model directory or Hugging Face repo id (required)
    GLIDE_MODEL_NAME          "model" field returned in responses [decision_config.json model_name, else glide-v1.2]
    GLIDE_DEVICE              torch device [cuda]
    GLIDE_MAX_LEN             max tokens per question prompt, image tokens included [40000]
    GLIDE_TOKEN_BUDGET        max padded tokens per forward micro-batch [16384]
    GLIDE_MAX_BATCH           max requests per micro-batch [16]
    GLIDE_BATCH_WAIT_MS       how long the worker waits to fill a micro-batch [5]
    GLIDE_MAX_QUEUE           max queued requests before HTTP 503 [1024]
    GLIDE_REQUEST_TIMEOUT_S   max time a request may wait for its result before HTTP 504 [600]
    GLIDE_THREADS             request handler threads [1024]
    GLIDE_API_KEY             if set, requests must send "Authorization: Bearer <key>"
    GLIDE_MIN_PIXELS / GLIDE_MAX_PIXELS   image resize bounds [65536 / 1048576]
"""
from __future__ import annotations

import argparse
import logging
import os
import queue
import threading
import time
from contextlib import asynccontextmanager

from .prompt import RequestError

log = logging.getLogger("glide.server")


def _env(name, default, cast=str):
    v = os.environ.get(name)
    return cast(v) if v not in (None, "") else default


class Settings:
    def __init__(self):
        self.model = _env("GLIDE_MODEL", "")
        self.name = _env("GLIDE_MODEL_NAME", None)          # unset: the name in decision_config.json (model_name), else glide-v1.2
        self.device = _env("GLIDE_DEVICE", "cuda")
        self.max_len = _env("GLIDE_MAX_LEN", 40000, int)
        self.token_budget = _env("GLIDE_TOKEN_BUDGET", 16384, int)
        self.max_batch = _env("GLIDE_MAX_BATCH", 16, int)
        self.batch_wait_ms = _env("GLIDE_BATCH_WAIT_MS", 5.0, float)
        self.max_queue = _env("GLIDE_MAX_QUEUE", 1024, int)
        self.request_timeout_s = _env("GLIDE_REQUEST_TIMEOUT_S", 600.0, float)
        self.threads = _env("GLIDE_THREADS", 1024, int)
        self.api_key = _env("GLIDE_API_KEY", "")


class Busy(Exception):
    pass


class _Job:
    __slots__ = ("request", "done", "result", "cancelled")

    def __init__(self, request):
        self.request, self.done, self.result, self.cancelled = request, threading.Event(), None, False


class MicroBatcher:
    """Collects concurrent requests for up to batch_wait_ms (at most max_batch of them) and scores them together."""

    def __init__(self, model, max_batch=16, batch_wait_ms=5.0, max_queue=1024):
        self.model = model
        self.max_batch, self.wait_s = max(1, int(max_batch)), max(0.0, float(batch_wait_ms)) / 1000
        self.queue = queue.Queue(maxsize=max(1, int(max_queue)))
        self.stats = {"requests": 0, "batches": 0, "busy_s": 0.0, "started": time.time()}
        threading.Thread(target=self._loop, name="glide-batcher", daemon=True).start()

    def submit(self, request, timeout):
        job = _Job(request)
        try:
            self.queue.put_nowait(job)
        except queue.Full:
            raise Busy() from None
        if not job.done.wait(timeout):
            job.cancelled = True
            raise TimeoutError()
        return job.result

    def _loop(self):
        while True:
            batch = [self.queue.get()]
            deadline = time.monotonic() + self.wait_s
            while len(batch) < self.max_batch:
                try:
                    batch.append(self.queue.get(timeout=max(0.0, deadline - time.monotonic())))
                except queue.Empty:
                    break
            batch = [j for j in batch if not j.cancelled]
            if not batch:
                continue
            t0 = time.perf_counter()
            results = self._run([j.request for j in batch])
            self.stats["busy_s"] += time.perf_counter() - t0
            self.stats["batches"] += 1
            self.stats["requests"] += len(batch)
            for j, r in zip(batch, results):
                j.result = r
                j.done.set()

    def _run(self, requests):
        try:
            return self.model.decide_batch(requests)
        except Exception as e:  # noqa: BLE001
            log.exception("batch of %d failed", len(requests))
            _empty_cache()
            if len(requests) == 1:
                return [e]
        out = []                      # isolate the failing request: score the batch one request at a time
        for r in requests:
            try:
                out.append(self.model.decide_batch([r])[0])
            except Exception as e:  # noqa: BLE001
                log.exception("request failed")
                _empty_cache()
                out.append(e)
        return out


def _empty_cache():
    try:
        import torch
        torch.cuda.empty_cache()
    except Exception:  # noqa: BLE001
        pass


def create_app(model=None, settings=None):
    """FastAPI app. model: an already loaded Glide instance; if None it is loaded from GLIDE_MODEL at startup."""
    from fastapi import FastAPI, Header
    from fastapi.responses import JSONResponse

    s = settings or Settings()
    ctx = {"model": model, "batcher": None}

    @asynccontextmanager
    async def lifespan(_app):
        import anyio.to_thread
        anyio.to_thread.current_default_thread_limiter().total_tokens = s.threads
        if ctx["model"] is None:
            if not s.model:
                raise RuntimeError("set GLIDE_MODEL to a model directory or Hugging Face repo id")
            from .model import Glide
            ctx["model"] = Glide.from_pretrained(s.model, device=s.device, max_len=s.max_len, token_budget=s.token_budget, name=s.name)
        ctx["batcher"] = MicroBatcher(ctx["model"], s.max_batch, s.batch_wait_ms, s.max_queue)
        yield

    app = FastAPI(title="Glide", lifespan=lifespan)

    def error(status, message):
        return JSONResponse(status_code=status, content={"error": message})

    @app.get("/health")
    def health():
        m, b = ctx["model"], ctx["batcher"]
        if m is None or b is None:
            return error(503, "loading")
        return {"status": "ok", "model": m.name, "vision": m.vision, "quantization": getattr(m, "quantization", None), "max_len": m.max_len, "token_budget": m.token_budget,
                "max_batch": b.max_batch, "queued": b.queue.qsize(), "stats": dict(b.stats)}

    @app.post("/v1/systemone")
    def systemone(body: dict, authorization: str = Header(default="")):
        if s.api_key and authorization != f"Bearer {s.api_key}":
            return error(401, "unauthorized")
        if "state" not in body:
            return error(422, "state is required")
        if not isinstance(body.get("questions"), dict):
            return error(422, "questions must be an object mapping keys to questions")
        t0 = time.perf_counter()
        try:
            r = ctx["batcher"].submit({"state": body["state"], "questions": body["questions"]}, s.request_timeout_s)
        except Busy:
            return error(503, "server busy: request queue is full")
        except TimeoutError:
            return error(504, "request timed out")
        if isinstance(r, RequestError):
            return error(422, str(r))
        if isinstance(r, Exception):
            return error(500, f"engine error: {type(r).__name__}")
        return {**r, "latency_ms": round((time.perf_counter() - t0) * 1000, 2)}

    return app


app = create_app()


def main():
    ap = argparse.ArgumentParser(description="Glide HTTP server (POST /v1/systemone, GET /health)")
    ap.add_argument("--model", default=os.environ.get("GLIDE_MODEL", ""), help="model directory or Hugging Face repo id")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=8000)
    args = ap.parse_args()
    if args.model:
        os.environ["GLIDE_MODEL"] = args.model
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    import uvicorn
    uvicorn.run(create_app(), host=args.host, port=args.port, workers=1)


if __name__ == "__main__":
    main()