"""Attempt-bound TLS result upload for Kubernetes Factory workers."""
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 = "Factory 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("FACTORY_K8S_RESULT_HOST must be an IPv4 address") from exc
        if address.version != 4 or not ATTEMPT_ID.fullmatch(attempt_id):
            raise ValueError("Factory result mailbox identity is invalid")
        self.advertised_host = advertised_host
        self.attempt_id = attempt_id
        self.output = output
        self.path = f"/v1/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="factory-result-mailbox-")
        self._server: _Server | None = None
        self._thread: threading.Thread | None = None
        self._ca_pem = ""

    def start(self) -> UploadTarget:
        if self._server is not None:
            raise RuntimeError("Factory result mailbox already started")
        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 Factory result certificate: {generated.stderr.strip()[-1200:]}")
        self._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"factory-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, self._ca_pem)

    def wait(self, timeout: float) -> Path:
        if not self.finished.wait(timeout):
            raise TimeoutError("Factory 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("Factory 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 __enter__(self) -> "ResultMailbox":
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        self.close()


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 target_sha256(target: UploadTarget) -> str:
    return hashlib.sha256(serialize_target(target).encode()).hexdigest()


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("Factory result upload URL is invalid") from exc
    target_attempt = parsed.path.removeprefix("/v1/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("Factory result upload URL is invalid")
    if attempt_id is not None and target_attempt != attempt_id:
        raise ValueError("Factory 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("Factory 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("Factory 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"Factory result upload failed with HTTP {response.status}")
    finally:
        connection.close()