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

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
from pathlib import Path
import re
import ssl
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)))
SPIFFE_PATH_SEGMENT = re.compile(r"^[A-Za-z0-9._-]+$")


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")
    if not job_name.strip():
        raise ValueError("identity components must not be empty")
    return "/".join(safe_identity_segment(part) for part in (*parts, job_name))


def safe_identity_segment(value: str) -> str:
    """Map arbitrary Gitea names to stable SPIFFE Operator path segments."""
    if SPIFFE_PATH_SEGMENT.fullmatch(value):
        return value
    slug = re.sub(r"[^A-Za-z0-9._-]+", "-", value).strip("-._")
    slug = slug[:48].rstrip("-._") or "segment"
    digest = hashlib.sha256(value.encode()).hexdigest()[:12]
    return f"{slug}-{digest}"


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]}"
    LOG.info(
        "runner request received job_id=%d run_id=%d repository=%s job_name=%r "
        "backend=%s pod=%s",
        request.job_id,
        request.run_id,
        request.repository,
        request.job_name,
        request.backend,
        pod_name,
    )
    stop = asyncio.Event()
    pulse = asyncio.create_task(heartbeat(message, stop))
    try:
        await client.create_pod(pod_manifest(request, pod_name))
        LOG.info(
            "runner Pod created pod=%s job_id=%d run_id=%d image=%s",
            pod_name,
            request.job_id,
            request.run_id,
            RUNNER_IMAGE,
        )
        if await wait_for_pod(client, pod_name):
            await message.ack()
            LOG.info(
                "runner request acknowledged pod=%s job_id=%d result=succeeded",
                pod_name,
                request.job_id,
            )
        else:
            LOG.error(
                "runner Pod failed pod=%s job_id=%d; request will be retried",
                pod_name,
                request.job_id,
            )
            await message.nak(delay=30)
    except Exception:
        LOG.exception(
            "runner request failed pod=%s job_id=%d; request will be retried",
            pod_name,
            request.job_id,
        )
        await message.nak(delay=30)
        raise
    finally:
        stop.set()
        await pulse
        LOG.info("deleting runner Pod pod=%s job_id=%d", pod_name, request.job_id)
        await client.delete_pod(pod_name)


async def bind_request(message: object, client: KubernetesClient) -> None:
    binding: IdentityBinding | None = 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)
        LOG.info(
            "identity binding received job_id=%d run_id=%d runner_name=%s "
            "repository=%s job_name=%r identity_path=%s",
            binding.job_id,
            binding.run_id,
            binding.runner_name,
            binding.repository,
            binding.job_name,
            path,
        )
        await client.bind_identity(binding.runner_name, path)
    except ClientResponseError as error:
        if error.status == 404:
            LOG.warning(
                "identity binding Pod not found runner_name=%s job_id=%s; binding will be retried",
                binding.runner_name if binding else None,
                binding.job_id if binding else None,
            )
            await message.nak(delay=2)
            return
        LOG.exception(
            "identity binding Kubernetes request failed runner_name=%s job_id=%s status=%d",
            binding.runner_name if binding else None,
            binding.job_id if binding else None,
            error.status,
        )
        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()
    LOG.info(
        "identity binding applied and acknowledged runner_name=%s job_id=%d identity_path=%s",
        binding.runner_name,
        binding.job_id,
        path,
    )


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()