#!/usr/bin/env python3
"""Capacity-bounded JetStream consumer that launches one ephemeral VM per job."""

import asyncio
import json
import logging
import os
import secrets
import ssl
import uuid
from pathlib import Path

import nats
from aiohttp import web
from nats.errors import TimeoutError
from nats.js.api import AckPolicy, ConsumerConfig

from .models import RunnerRequest

LOG = logging.getLogger(__name__)
CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "1"))
SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.vm")
STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
DURABLE = os.environ.get("NATS_DURABLE", "vm")
MAX_INFLIGHT = int(os.environ.get("RUNNER_MAX_INFLIGHT", "64"))
NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
NATS_USER = os.environ.get("NATS_USER", "ci-worker")
NATS_PASSWORD_FILE = Path(os.environ.get("NATS_PASSWORD_FILE", "/etc/microvm-runner/nats-password"))
NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
REGISTRATION_TOKEN_FILE = Path(os.environ.get("REGISTRATION_TOKEN_FILE", "/etc/microvm-runner/registration-token"))
LAUNCHER = os.environ.get("LAUNCHER", "/usr/local/libexec/microvm-runner-launch")
TOKEN_LISTEN = os.environ.get("TOKEN_LISTEN", "172.30.0.1")
TOKEN_PORT = int(os.environ.get("TOKEN_PORT", "8787"))
GUEST_ASSETS = Path(
    os.environ.get("GUEST_ASSETS", "/opt/gitea-dynamic-runner/guest-assets.tar.gz")
)

tokens: dict[str, bytes] = {}
token_lock = asyncio.Lock()


async def token(request: web.Request) -> web.Response:
    nonce = request.match_info["nonce"]
    async with token_lock:
        value = tokens.pop(nonce, None)
    if value is None:
        raise web.HTTPNotFound()
    return web.Response(body=value, headers={"Cache-Control": "no-store"})


async def guest_assets(_: web.Request) -> web.FileResponse:
    return web.FileResponse(
        GUEST_ASSETS,
        headers={"Cache-Control": "public, immutable"},
    )


async def heartbeat(message: object, stop: asyncio.Event) -> None:
    while True:
        try:
            await asyncio.wait_for(stop.wait(), timeout=60)
            return
        except asyncio.TimeoutError:
            await message.in_progress()


async def run_one(message: object) -> None:
    try:
        request = RunnerRequest.from_json(message.data)
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as error:
        LOG.error("discarding invalid runner request: %s", error)
        await message.ack()
        return
    if request.backend != "vm":
        LOG.error("discarding %s request received by VM worker", request.backend)
        await message.ack()
        return

    instance_id = str(uuid.uuid4())
    nonce = secrets.token_urlsafe(32)
    async with token_lock:
        tokens[nonce] = REGISTRATION_TOKEN_FILE.read_bytes().strip()
    stop = asyncio.Event()
    pulse = asyncio.create_task(heartbeat(message, stop))
    try:
        process = await asyncio.create_subprocess_exec(
            LAUNCHER,
            instance_id,
            nonce,
            env={
                **os.environ,
                "RUNNER_JOB_ID": str(request.job_id),
                "RUNNER_RUN_ID": str(request.run_id),
                "RUNNER_REPOSITORY": request.repository,
                "RUNNER_JOB_NAME": request.job_name,
            },
        )
        return_code = await process.wait()
        if return_code == 0:
            await message.ack()
        else:
            LOG.error("launcher for %s exited with %d", instance_id, return_code)
            await message.nak(delay=30)
    except Exception:
        await message.nak(delay=30)
        raise
    finally:
        stop.set()
        await pulse
        async with token_lock:
            tokens.pop(nonce, None)


def _report_task(task: asyncio.Task[None]) -> None:
    if task.cancelled():
        return
    error = task.exception()
    if error is not None:
        LOG.error(
            "runner task failed",
            exc_info=(type(error), error, error.__traceback__),
        )


async def consume() -> None:
    if CAPACITY < 1:
        raise ValueError("RUNNER_CAPACITY must be at least 1")
    tls = ssl.create_default_context(cafile=NATS_CA_FILE)
    nc = await nats.connect(
        NATS_URL,
        user=NATS_USER,
        password=NATS_PASSWORD_FILE.read_text().strip(),
        tls=tls,
        name=DURABLE,
    )
    js = nc.jetstream()
    subscription = await js.pull_subscribe(
        SUBJECT,
        durable=DURABLE,
        stream=STREAM,
        config=ConsumerConfig(
            durable_name=DURABLE,
            filter_subject=SUBJECT,
            ack_policy=AckPolicy.EXPLICIT,
            ack_wait=5 * 60,
            max_ack_pending=MAX_INFLIGHT,
            max_deliver=5,
        ),
    )
    active: set[asyncio.Task[None]] = set()
    try:
        while True:
            active = {task for task in active if not task.done()}
            free = CAPACITY - len(active)
            if free == 0:
                await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED)
                continue
            try:
                messages = await subscription.fetch(batch=free, timeout=5)
            except TimeoutError:
                continue
            for message in messages:
                task = asyncio.create_task(run_one(message))
                task.add_done_callback(_report_task)
                active.add(task)
    finally:
        if active:
            await asyncio.gather(*active, return_exceptions=True)
        await nc.drain()


async def main() -> None:
    app = web.Application()
    app.router.add_get("/token/{nonce}", token)
    app.router.add_get("/assets/guest-assets.tar.gz", guest_assets)
    runner = web.AppRunner(app)
    await runner.setup()
    await web.TCPSite(runner, TOKEN_LISTEN, TOKEN_PORT).start()
    try:
        await consume()
    finally:
        await runner.cleanup()


def cli() -> None:
    logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
    asyncio.run(main())


if __name__ == "__main__":
    cli()