#!/usr/bin/env python3
"""JetStream worker that creates one disposable Kubernetes Pod per job."""

from __future__ import annotations

import asyncio
import json
import logging
import os
from pathlib import Path
import ssl
from urllib.parse import quote
import uuid

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

from .kubernetes import KubernetesClient
from .models import IdentityBinding, RunnerRequest


LOG = logging.getLogger(__name__)
STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
REQUEST_SUBJECT = os.environ.get("NATS_SUBJECT", "ci.runner.pod")
BINDING_SUBJECT = os.environ.get("NATS_BINDING_SUBJECT", "ci.runner.pod.binding")
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", "/run/secrets/nats/password"))
NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
KUBERNETES_API = os.environ.get("KUBERNETES_API", "https://kubernetes.default.svc")
KUBERNETES_TOKEN_FILE = Path(
    os.environ.get(
        "KUBERNETES_TOKEN_FILE",
        "/var/run/secrets/kubernetes.io/serviceaccount/token",
    )
)
KUBERNETES_CA_FILE = Path(
    os.environ.get(
        "KUBERNETES_CA_FILE",
        "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt",
    )
)
NAMESPACE = os.environ.get("RUNNER_NAMESPACE", "gitea-actions")
RUNNER_IMAGE = os.environ.get("RUNNER_IMAGE", "docker.io/gitea/runner:2")
RUNNER_SERVICE_ACCOUNT = os.environ.get("RUNNER_SERVICE_ACCOUNT", "gitea-dynamic-runner")
RUNNER_TOKEN_SECRET = os.environ.get("RUNNER_TOKEN_SECRET", "gitea-dynamic-runner")
GITEA_INSTANCE = os.environ.get("GITEA_INSTANCE", "https://git.ddupan.top")
CAPACITY = int(os.environ.get("RUNNER_CAPACITY", "4"))
POD_TIMEOUT = int(os.environ.get("RUNNER_POD_TIMEOUT", str(4 * 60 * 60)))


def identity_path(repository: str, job_name: str) -> str:
    parts = repository.split("/")
    if len(parts) != 2 or not all(parts):
        raise ValueError("repository must be owner/name")
    encoded = [quote(part, safe="-._~") for part in (*parts, job_name)]
    if not job_name.strip() or any(not part for part in encoded):
        raise ValueError("identity components must not be empty")
    return "/".join(encoded)


def pod_manifest(request: RunnerRequest, pod_name: str) -> dict[str, object]:
    if request.backend != "pod":
        raise ValueError("Pod backend only accepts pod requests")
    return {
        "apiVersion": "v1",
        "kind": "Pod",
        "metadata": {
            "name": pod_name,
            "namespace": NAMESPACE,
            "labels": {
                "app.kubernetes.io/name": "gitea-dynamic-runner",
                "app.kubernetes.io/component": "runner",
                "ci.ddupan.top/backend": "pod",
            },
            "annotations": {
                "ci.ddupan.top/queued-job-id": str(request.job_id),
                "ci.ddupan.top/queued-run-id": str(request.run_id),
            },
        },
        "spec": {
            "serviceAccountName": RUNNER_SERVICE_ACCOUNT,
            "restartPolicy": "Never",
            "terminationGracePeriodSeconds": 30,
            "containers": [
                {
                    "name": "runner",
                    "image": RUNNER_IMAGE,
                    "imagePullPolicy": "IfNotPresent",
                    "securityContext": {"privileged": True},
                    "env": [
                        {"name": "GITEA_INSTANCE_URL", "value": GITEA_INSTANCE},
                        {
                            "name": "GITEA_RUNNER_NAME",
                            "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}},
                        },
                        {
                            "name": "GITEA_RUNNER_REGISTRATION_TOKEN_FILE",
                            "value": "/run/secrets/gitea/token",
                        },
                        {"name": "GITEA_RUNNER_LABELS", "value": "self-hosted:host,pod:host"},
                        {"name": "GITEA_RUNNER_EPHEMERAL", "value": "1"},
                        {"name": "GITEA_RUNNER_ONCE", "value": "1"},
                        {"name": "CONFIG_FILE", "value": "/etc/gitea-runner/config.yaml"},
                    ],
                    "volumeMounts": [
                        {
                            "name": "registration-token",
                            "mountPath": "/run/secrets/gitea",
                            "readOnly": True,
                        },
                        {
                            "name": "spire-agent-socket",
                            "mountPath": "/run/spire/agent-sockets",
                            "readOnly": True,
                        },
                    ],
                }
            ],
            "volumes": [
                {
                    "name": "registration-token",
                    "secret": {
                        "secretName": RUNNER_TOKEN_SECRET,
                        "items": [{"key": "token", "path": "token"}],
                    },
                },
                {
                    "name": "spire-agent-socket",
                    "csi": {"driver": "csi.spiffe.io", "readOnly": True},
                },
            ],
        },
    }


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 wait_for_pod(client: KubernetesClient, name: str) -> bool:
    deadline = asyncio.get_running_loop().time() + POD_TIMEOUT
    while asyncio.get_running_loop().time() < deadline:
        pod = await client.get_pod(name)
        if pod is None:
            raise RuntimeError(f"runner Pod {name} disappeared")
        status = pod.get("status")
        phase = status.get("phase") if isinstance(status, dict) else None
        if phase == "Succeeded":
            return True
        if phase == "Failed":
            return False
        await asyncio.sleep(2)
    raise asyncio.TimeoutError(f"runner Pod {name} timed out")


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

    pod_name = f"gitea-pod-{uuid.uuid4().hex[:12]}"
    stop = asyncio.Event()
    pulse = asyncio.create_task(heartbeat(message, stop))
    try:
        await client.create_pod(pod_manifest(request, pod_name))
        if await wait_for_pod(client, pod_name):
            await message.ack()
        else:
            LOG.error("runner Pod %s failed", pod_name)
            await message.nak(delay=30)
    except Exception:
        await message.nak(delay=30)
        raise
    finally:
        stop.set()
        await pulse
        await client.delete_pod(pod_name)


async def bind_request(message: object, client: KubernetesClient) -> None:
    try:
        document = json.loads(message.data)
        binding = IdentityBinding(**document)
        if binding.backend != "pod" or not binding.runner_name.startswith("gitea-pod-"):
            raise ValueError("invalid Pod identity binding")
        path = identity_path(binding.repository, binding.job_name)
        await client.bind_identity(binding.runner_name, path)
    except ClientResponseError as error:
        if error.status == 404:
            await message.nak(delay=2)
            return
        await message.nak(delay=30)
        raise
    except (json.JSONDecodeError, TypeError, ValueError) as error:
        LOG.error("discarding invalid identity binding: %s", error)
        await message.ack()
        return
    await message.ack()


async def consume_requests(subscription: object, client: KubernetesClient) -> None:
    active: set[asyncio.Task[None]] = set()
    while True:
        active = {task for task in active if not task.done()}
        free = CAPACITY - len(active)
        if free < 1:
            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_request(message, client))
            task.add_done_callback(_report_task)
            active.add(task)


async def consume_bindings(subscription: object, client: KubernetesClient) -> None:
    while True:
        try:
            messages = await subscription.fetch(batch=16, timeout=5)
        except TimeoutError:
            continue
        await asyncio.gather(*(bind_request(message, client) for message in messages))


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


async def main() -> None:
    if CAPACITY < 1:
        raise ValueError("RUNNER_CAPACITY must be at least 1")
    context = 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=context,
        name="pod-runner-worker",
    )
    js = nc.jetstream()
    requests = await js.pull_subscribe(
        REQUEST_SUBJECT,
        durable="pod",
        stream=STREAM,
        config=ConsumerConfig(
            durable_name="pod",
            filter_subject=REQUEST_SUBJECT,
            ack_policy=AckPolicy.EXPLICIT,
            ack_wait=5 * 60,
            max_ack_pending=max(CAPACITY, 1),
            max_deliver=5,
        ),
    )
    bindings = await js.pull_subscribe(
        BINDING_SUBJECT,
        durable="pod-binding",
        stream=STREAM,
        config=ConsumerConfig(
            durable_name="pod-binding",
            filter_subject=BINDING_SUBJECT,
            ack_policy=AckPolicy.EXPLICIT,
            ack_wait=30,
            max_ack_pending=64,
            max_deliver=10,
        ),
    )
    async with KubernetesClient(
        api_url=KUBERNETES_API,
        token_file=KUBERNETES_TOKEN_FILE,
        ca_file=KUBERNETES_CA_FILE,
        namespace=NAMESPACE,
    ) as client:
        try:
            await asyncio.gather(
                consume_requests(requests, client),
                consume_bindings(bindings, client),
            )
        finally:
            await nc.drain()


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


if __name__ == "__main__":
    cli()