"""QR Scan Difficulty Check: the measured bell, and a place on it for your own pictures.

GET  /                 the page (web/)
POST /api/measure      one picture (JPEG, PNG or WebP, at most 25 MB) as the request body; returns its reads by
                       the five scanners at the seven frame sizes and its place on the scale
GET  /api/health       what is running
Pictures are read from a temporary file that is deleted when the reading ends; nothing is stored, and no access log is
kept.
"""

import asyncio
import os
import tempfile
from contextlib import asynccontextmanager
from importlib.metadata import version
from pathlib import Path

from starlette.applications import Starlette
from starlette.concurrency import run_in_threadpool
from starlette.requests import ClientDisconnect, Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route
from starlette.staticfiles import StaticFiles

from .place import place
from .reading import FRAMES, SCANNERS, PictureError, Reader

WEB = Path(__file__).resolve().parent.parent / "web"
MAX_BYTES = 25 * 1024 * 1024
SLOTS = int(os.environ.get("QSD_SLOTS", "2"))   # pictures read at once: one decoder worker each
MAX_WAITING = 16      # pictures waiting beyond those (arriving, or waiting for a slot); more are turned away
SLOT_WAIT_S = 10      # the longest a picture waits for a slot, so that a request ends within a minute
BODY_S = 60           # the longest a picture may take to arrive; a stalled upload doesn't keep its place
BUSY = "Many pictures are being read right now. Try again in a minute."
CSP = "; ".join([
    "default-src 'self'", "img-src 'self' blob: data:", "style-src 'self' 'unsafe-inline'", "script-src 'self'",
    "connect-src 'self'", "object-src 'none'", "base-uri 'none'", "form-action 'none'",
    "frame-ancestors 'self' https://huggingface.co",
])
STATE: dict = {}


def sniff(data: bytes) -> str | None:
    """The picture's format from its first bytes: only JPEG, PNG and WebP are opened here (the page converts others)."""
    if data[:3] == b"\xff\xd8\xff":
        return ".jpg"
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return ".png"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return ".webp"
    return None


def problem(status: int, message: str) -> JSONResponse:
    return JSONResponse({"error": message}, status_code=status, headers={"Cache-Control": "no-store"})


async def body(request: Request) -> bytes | None:
    """The request body, or None past MAX_BYTES (read no further)."""
    n = request.headers.get("content-length", "")
    if n.isdigit() and int(n) > MAX_BYTES:
        return None
    buf = bytearray()   # one buffer, so a picture costs its own size in memory however small its chunks
    async for chunk in request.stream():
        buf += chunk
        if len(buf) > MAX_BYTES:
            return None
    return bytes(buf)


def read_picture(data: bytes, suffix: str) -> dict:
    with tempfile.NamedTemporaryFile(prefix="qsd-", suffix=suffix) as f:
        f.write(data)
        f.flush()
        result = STATE["reader"].read(f.name)
    result["order"] = FRAMES   # the frame size behind each position of "reads" and "frames"
    result["place"] = place(result["reads"], FRAMES)
    result["count1280"] = sum(result["reads"][k][FRAMES.index(1280)] for k in SCANNERS)
    return result


async def measure(request: Request) -> JSONResponse:
    if STATE["busy"] >= SLOTS + MAX_WAITING:
        return problem(503, BUSY)
    STATE["busy"] += 1   # counted from the first byte, so at most SLOTS + MAX_WAITING pictures are held in memory
    try:
        try:
            async with asyncio.timeout(BODY_S):
                data = await body(request)
        except TimeoutError:
            return problem(408, "This picture took too long to arrive. Try again, or send a smaller copy.")
        except ClientDisconnect:   # the visitor left mid-upload: nobody will read this answer
            return problem(400, "The picture did not arrive.")
        if data is None:
            return problem(413, "This picture is over 25 MB.")
        suffix = sniff(data)
        if suffix is None:
            return problem(415, "Send a JPEG, PNG or WebP picture.")
        try:
            await asyncio.wait_for(STATE["gate"].acquire(), SLOT_WAIT_S)
        except TimeoutError:
            return problem(503, BUSY)
        try:
            result = await run_in_threadpool(read_picture, data, suffix)
        finally:
            STATE["gate"].release()
    except PictureError as e:
        return problem(422, str(e))
    finally:
        STATE["busy"] -= 1
    return JSONResponse(result, headers={"Cache-Control": "no-store"})


async def health(request: Request) -> JSONResponse:
    return JSONResponse({
        "ok": True, "scanners": SCANNERS, "frames": FRAMES, "slots": SLOTS,
        "versions": {"opencv-contrib": version("opencv-contrib-python-headless"), "zxing-cpp": version("zxing-cpp"),
                     "zxing": "3.5.2"},
    }, headers={"Cache-Control": "no-store"})


@asynccontextmanager
async def lifespan(app):
    STATE["reader"] = await run_in_threadpool(Reader, SLOTS)
    STATE["gate"] = asyncio.Semaphore(SLOTS)
    STATE["busy"] = 0
    yield
    STATE["reader"].close()


class SecurityHeaders:
    """Plain ASGI middleware: the same headers on every response."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        async def send_with_headers(message):
            if message["type"] == "http.response.start":
                message.setdefault("headers", [])
                message["headers"] += [(b"content-security-policy", CSP.encode()),
                                       (b"x-content-type-options", b"nosniff"),
                                       (b"referrer-policy", b"strict-origin-when-cross-origin")]
            await send(message)

        return await self.app(scope, receive, send_with_headers)


app = SecurityHeaders(Starlette(lifespan=lifespan, routes=[
    Route("/api/measure", measure, methods=["POST"]),
    Route("/api/health", health, methods=["GET"]),
    Mount("/", StaticFiles(directory=WEB, html=True), name="web"),
]))