"""Bounded, attempt-scoped result transport for cdx Kubernetes Jobs."""
from __future__ import annotations

import hashlib
import hmac
import http.client
import ipaddress
import json
import os
import re
import secrets
import ssl
import subprocess
import tempfile
import threading
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlsplit

MAX_BUNDLE_BYTES = 256 * 1024 * 1024
MAX_CONNECTIONS = 8
REQUEST_TIMEOUT_SECONDS = 30
ATTEMPT_ID = re.compile(r"^[0-9a-f]{24}$")
SHA256 = re.compile(r"^[0-9a-f]{64}$")
TOKEN = re.compile(r"^[A-Za-z0-9_-]{43,128}$")


@dataclass(frozen=True)
class UploadTarget:
    url: str
    token: str
    ca_pem: str


class _Handler(BaseHTTPRequestHandler):
    server: "_Server"

    def do_PUT(self) -> None:
        mailbox = self.server.mailbox
        if self.path != mailbox.path:
            self.send_error(404)
            return
        authorization = self.headers.get("Authorization", "")
        if not hmac.compare_digest(authorization, f"Bearer {mailbox.token}"):
            self.send_error(401)
            return
        if self.headers.get("Content-Type") != "application/octet-stream" or self.headers.get("Transfer-Encoding"):
            self.send_error(415)
            return
        try:
            length = int(self.headers.get("Content-Length", ""))
        except ValueError:
            self.send_error(400)
            return
        expected_hash = self.headers.get("X-Content-SHA256", "")
        if not 0 < length <= MAX_BUNDLE_BYTES or not SHA256.fullmatch(expected_hash):
            self.send_error(400)
            return
        with mailbox.lock:
            if mailbox.claimed:
                self.send_error(409)
                return
            mailbox.claimed = True
        digest = hashlib.sha256()
        remaining = length
        try:
            with mailbox.temporary_path.open("xb") as output:
                while remaining:
                    chunk = self.rfile.read(min(1024 * 1024, remaining))
                    if not chunk:
                        raise OSError("result upload ended before Content-Length")
                    output.write(chunk)
                    digest.update(chunk)
                    remaining -= len(chunk)
            if not hmac.compare_digest(digest.hexdigest(), expected_hash):
                raise ValueError("result upload digest mismatch")
            os.replace(mailbox.temporary_path, mailbox.bundle_path)
        except (OSError, ValueError):
            mailbox.temporary_path.unlink(missing_ok=True)
            mailbox.failure = "cdx result upload failed validation"
            mailbox.finished.set()
            try:
                self.send_error(400)
            except (OSError, ssl.SSLError):
                pass
            return
        mailbox.ready.set()
        mailbox.finished.set()
        self.send_response(201)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, format: str, *args: object) -> None:
        return


class _Server(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, address: tuple[str, int], mailbox: "ResultMailbox", context: ssl.SSLContext):
        self.mailbox = mailbox
        self.context = context
        self.connections = threading.BoundedSemaphore(MAX_CONNECTIONS)
        super().__init__(address, _Handler)

    def get_request(self):
        connection, address = super().get_request()
        connection.settimeout(REQUEST_TIMEOUT_SECONDS)
        try:
            secured = self.context.wrap_socket(connection, server_side=True, do_handshake_on_connect=False)
        except BaseException:
            connection.close()
            raise
        return secured, address

    def process_request(self, request, client_address) -> None:
        if not self.connections.acquire(blocking=False):
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except BaseException:
            self.connections.release()
            raise

    def process_request_thread(self, request, client_address) -> None:
        try:
            super().process_request_thread(request, client_address)
        finally:
            self.connections.release()


class ResultMailbox:
    def __init__(self, advertised_host: str, attempt_id: str, output: Path):
        try:
            address = ipaddress.ip_address(advertised_host)
        except ValueError as exc:
            raise ValueError("cdx result host must be an IPv4 address") from exc
        if address.version != 4 or not ATTEMPT_ID.fullmatch(attempt_id):
            raise ValueError("cdx result mailbox identity is invalid")
        self.advertised_host = advertised_host
        self.attempt_id = attempt_id
        self.output = output
        self.path = f"/v1/cdx-results/{attempt_id}"
        self.token = secrets.token_urlsafe(32)
        self.bundle_path = output / "result.bundle"
        self.temporary_path = output / "result.bundle.uploading"
        self.lock = threading.Lock()
        self.ready = threading.Event()
        self.finished = threading.Event()
        self.failure: str | None = None
        self.claimed = False
        self._temporary = tempfile.TemporaryDirectory(prefix="cdx-result-mailbox-")
        self._server: _Server | None = None
        self._thread: threading.Thread | None = None

    def start(self) -> UploadTarget:
        if self._server is not None:
            raise RuntimeError("cdx result mailbox already started")
        self.output.mkdir(parents=True, exist_ok=True)
        directory = Path(self._temporary.name)
        certificate = directory / "certificate.pem"
        private_key = directory / "private-key.pem"
        generated = subprocess.run([
            "openssl", "req", "-x509", "-newkey", "rsa:2048", "-sha256", "-nodes",
            "-keyout", str(private_key), "-out", str(certificate), "-days", "1",
            "-subj", f"/CN={self.advertised_host}", "-addext", f"subjectAltName=IP:{self.advertised_host}",
        ], text=True, capture_output=True, check=False, timeout=30)
        if generated.returncode:
            raise RuntimeError(f"could not create cdx result certificate: {generated.stderr.strip()[-1200:]}")
        ca_pem = certificate.read_text()
        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        context.load_cert_chain(certificate, private_key)
        server = _Server(("0.0.0.0", 0), self, context)
        self._server = server
        self._thread = threading.Thread(target=server.serve_forever, name=f"cdx-result-{self.attempt_id}", daemon=True)
        self._thread.start()
        port = server.server_address[1]
        return UploadTarget(f"https://{self.advertised_host}:{port}{self.path}", self.token, ca_pem)

    def wait(self, timeout: float) -> Path:
        if not self.finished.wait(timeout):
            raise TimeoutError("cdx worker did not upload a result bundle")
        if self.failure is not None:
            raise RuntimeError(self.failure)
        if not self.ready.is_set() or not self.bundle_path.is_file() or not self.bundle_path.stat().st_size:
            raise RuntimeError("cdx result mailbox produced an empty bundle")
        return self.bundle_path

    def close(self) -> None:
        if self._server is not None:
            self._server.shutdown()
            self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=5)
        try:
            self.temporary_path.unlink(missing_ok=True)
        finally:
            try:
                self.bundle_path.unlink(missing_ok=True)
            finally:
                self._temporary.cleanup()


def serialize_target(target: UploadTarget) -> str:
    validate_target(target)
    return json.dumps({"url": target.url, "token": target.token, "ca_pem": target.ca_pem}, sort_keys=True, separators=(",", ":"))


def validate_target(target: UploadTarget, attempt_id: str | None = None):
    try:
        parsed = urlsplit(target.url)
        port = parsed.port
        address = ipaddress.ip_address(parsed.hostname or "")
    except ValueError as exc:
        raise ValueError("cdx result upload URL is invalid") from exc
    target_attempt = parsed.path.removeprefix("/v1/cdx-results/")
    if parsed.scheme != "https" or parsed.username or parsed.password or parsed.query or parsed.fragment or address.version != 4 or port is None or not ATTEMPT_ID.fullmatch(target_attempt):
        raise ValueError("cdx result upload URL is invalid")
    if attempt_id is not None and target_attempt != attempt_id:
        raise ValueError("cdx result upload attempt does not match execution")
    if not TOKEN.fullmatch(target.token) or not target.ca_pem.startswith("-----BEGIN CERTIFICATE-----") or not target.ca_pem.rstrip().endswith("-----END CERTIFICATE-----"):
        raise ValueError("cdx result upload authorization is invalid")
    return parsed


def upload_result(target: UploadTarget, bundle: Path, *, timeout: float = 120) -> None:
    parsed = validate_target(target)
    size = bundle.stat().st_size
    if not 0 < size <= MAX_BUNDLE_BYTES:
        raise ValueError("cdx result bundle size is invalid")
    digest = hashlib.sha256()
    with bundle.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    context = ssl.create_default_context(cadata=target.ca_pem)
    connection = http.client.HTTPSConnection(parsed.hostname, parsed.port, timeout=timeout, context=context)
    try:
        connection.putrequest("PUT", parsed.path)
        connection.putheader("Authorization", f"Bearer {target.token}")
        connection.putheader("Content-Type", "application/octet-stream")
        connection.putheader("Content-Length", str(size))
        connection.putheader("X-Content-SHA256", digest.hexdigest())
        connection.endheaders()
        with bundle.open("rb") as source:
            for chunk in iter(lambda: source.read(1024 * 1024), b""):
                connection.send(chunk)
        response = connection.getresponse()
        response.read(1024)
        if response.status != 201:
            raise RuntimeError(f"cdx result upload failed with HTTP {response.status}")
    finally:
        connection.close()