#!/usr/bin/env python3
"""Deterministic verification and publication helpers for the release image."""

from __future__ import annotations

import argparse
import hashlib
import os
import re
import subprocess
import sys
import tarfile
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol, Sequence

IMAGE = "ghcr.io/fourblocks2020/trex-core"
CANARY = b"TREX_SECRET_CANARY_DO_NOT_SHIP"
COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
DIGEST_RE = re.compile(r"(?:^|\s)digest:\s*(sha256:[0-9a-f]{64})(?:\s|$)")
IMAGE_ID_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
MAX_ARCHIVE_DEPTH = 8
MAX_ARCHIVE_MEMBERS = 200_000
MAX_ARCHIVE_EXTRACTED_BYTES = 2 * 1024 * 1024 * 1024
MAX_ARCHIVE_MEMBER_BYTES = 512 * 1024 * 1024
MAX_ARCHIVE_INPUT_BYTES = 1024 * 1024 * 1024


@dataclass(frozen=True)
class Result:
    stdout: str = ""
    stderr: str = ""


@dataclass
class _ScanBudget:
    max_depth: int = MAX_ARCHIVE_DEPTH
    max_members: int = MAX_ARCHIVE_MEMBERS
    max_extracted_bytes: int = MAX_ARCHIVE_EXTRACTED_BYTES
    max_member_bytes: int = MAX_ARCHIVE_MEMBER_BYTES
    max_input_bytes: int = MAX_ARCHIVE_INPUT_BYTES
    members: int = 0
    extracted_bytes: int = 0

    def check_depth(self, depth: int) -> None:
        if depth > self.max_depth:
            _scan_limit_exceeded()

    def add_member(self, size: int) -> None:
        self.members += 1
        if self.members > self.max_members:
            _scan_limit_exceeded()
        if size < 0 or size > self.max_member_bytes:
            _scan_limit_exceeded()
        self.extracted_bytes += size
        if self.extracted_bytes > self.max_extracted_bytes:
            _scan_limit_exceeded()


def _scan_limit_exceeded() -> None:
    raise ValueError("archive scan resource limit exceeded")


class Runner(Protocol):
    def run(
        self,
        command: Sequence[str],
        *,
        input_text: str | None = None,
        timeout: int | None = None,
    ) -> Result: ...


class SubprocessRunner:
    def run(
        self,
        command: Sequence[str],
        *,
        input_text: str | None = None,
        timeout: int | None = None,
    ) -> Result:
        completed = subprocess.run(
            command,
            input=input_text,
            text=True,
            capture_output=True,
            check=True,
            timeout=timeout,
        )
        return Result(completed.stdout, completed.stderr)


def validate_commit(commit: str) -> str:
    if not COMMIT_RE.fullmatch(commit):
        raise ValueError("source commit must be 40 lowercase hexadecimal characters")
    return commit


def release_tags(commit: str, ref: str) -> list[str]:
    validate_commit(commit)
    tags = [f"{IMAGE}:{commit[:12]}"]
    if ref == "refs/heads/main":
        return tags
    prefix = "refs/tags/"
    if ref.startswith(prefix) and ref.removeprefix(prefix).startswith("v"):
        release = ref.removeprefix(prefix)
        if not re.fullmatch(r"v[A-Za-z0-9._-]*", release):
            raise ValueError("release ref contains unsupported tag characters")
        tags.append(f"{IMAGE}:{release}")
        return tags
    raise ValueError("publication ref must be main or a v* tag")


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_checksum(archive: Path, checksum: Path) -> None:
    checksum.write_text(f"{sha256_file(archive)}  {archive.name}\n", encoding="ascii")


def verify_checksum(archive: Path, checksum: Path) -> None:
    if not archive.is_file() or not checksum.is_file():
        raise ValueError("image archive or checksum evidence is missing")
    fields = checksum.read_text(encoding="ascii").strip().split()
    if len(fields) != 2 or fields[1] != archive.name or not re.fullmatch(r"[0-9a-f]{64}", fields[0]):
        raise ValueError("checksum evidence has an invalid format")
    if sha256_file(archive) != fields[0]:
        raise ValueError("image archive checksum mismatch")


def validate_image_id(image_id: str) -> str:
    if not IMAGE_ID_RE.fullmatch(image_id):
        raise ValueError("image ID must be sha256 followed by 64 lowercase hexadecimal characters")
    return image_id


def inspect_image_id(image: str, runner: Runner) -> str:
    result = runner.run(["docker", "image", "inspect", "--format", "{{.Id}}", image])
    return validate_image_id(result.stdout.strip())


def write_image_id(image: str, output: Path, runner: Runner) -> str:
    image_id = inspect_image_id(image, runner)
    output.write_text(f"{image_id}\n", encoding="ascii")
    return image_id


def read_image_id(path: Path) -> str:
    text = path.read_text(encoding="ascii")
    if not re.fullmatch(r"sha256:[0-9a-f]{64}\n?", text):
        raise ValueError("tested image ID evidence has an invalid format")
    return text.rstrip("\n")


def verify_image_id(image: str, expected_path: Path, runner: Runner) -> str:
    expected = read_image_id(expected_path)
    actual = inspect_image_id(image, runner)
    if actual != expected:
        raise ValueError("loaded image ID does not match tested image ID")
    return actual


SMOKE_PROBES: tuple[tuple[str, ...], ...] = (
    ("python", "-c", "import trex_core"),
    ("trex-core-daemon", "--help"),
    ("trex-core", "--help"),
    ("trex-core-reset", "--help"),
)


def _ensure_result_clean(result: Result) -> None:
    if CANARY in (result.stdout + result.stderr).encode():
        raise ValueError("canary marker found on a scanned surface")


def smoke_image(image: str, runner: Runner, timeout: int = 30) -> None:
    for index, probe in enumerate(SMOKE_PROBES):
        name = f"trex-core-smoke-{os.getpid()}-{index}"
        command = [
            "docker",
            "run",
            "--name",
            name,
            "--network",
            "none",
            "--read-only",
            "--cap-drop",
            "ALL",
            "--security-opt",
            "no-new-privileges",
            "--pids-limit",
            "64",
            "--tmpfs",
            "/tmp:rw,noexec,nosuid,size=16m",
            "--entrypoint",
            probe[0],
            image,
            *probe[1:],
        ]
        failure: BaseException | None = None
        try:
            _ensure_result_clean(runner.run(command, timeout=timeout))
        except BaseException as error:
            failure = error
        try:
            _ensure_result_clean(
                runner.run(["docker", "rm", "--force", name], timeout=timeout)
            )
        except BaseException as cleanup_error:
            if failure is None:
                raise RuntimeError("smoke container cleanup failed") from cleanup_error
        if failure is not None:
            raise RuntimeError(f"smoke probe {index + 1} failed") from failure


def _scan_stream(path: Path) -> None:
    with path.open("rb") as stream:
        overlap = b""
        while chunk := stream.read(1024 * 1024):
            data = overlap + chunk
            if CANARY in data:
                raise ValueError("canary marker found on a scanned surface")
            overlap = data[-(len(CANARY) - 1):]


def _scan_metadata(member: tarfile.TarInfo) -> None:
    values = [member.name, member.linkname, member.uname, member.gname]
    for key, value in member.pax_headers.items():
        values.extend((key, value))
    if any(CANARY in str(value).encode("utf-8", errors="replace") for value in values):
        raise ValueError("canary marker found on a scanned surface")


def _scan_tar(path: Path, budget: _ScanBudget | None = None, depth: int = 0) -> None:
    budget = budget or _ScanBudget()
    if path.stat().st_size > (budget.max_input_bytes if depth == 0 else budget.max_member_bytes):
        _scan_limit_exceeded()
    _scan_stream(path)
    try:
        archive = tarfile.open(path, mode="r:*")
    except (tarfile.TarError, OSError, EOFError):
        return
    with archive:
        budget.check_depth(depth)
        for member in archive:
            budget.add_member(member.size if member.isfile() else 0)
            _scan_metadata(member)
            if not member.isfile():
                continue
            extracted = archive.extractfile(member)
            if extracted is None:
                continue
            with tempfile.NamedTemporaryFile() as nested:
                overlap = b""
                copied = 0
                while chunk := extracted.read(1024 * 1024):
                    copied += len(chunk)
                    if copied > member.size or copied > budget.max_member_bytes:
                        _scan_limit_exceeded()
                    data = overlap + chunk
                    if CANARY in data:
                        raise ValueError("canary marker found on a scanned surface")
                    nested.write(chunk)
                    overlap = data[-(len(CANARY) - 1):]
                nested.flush()
                _scan_tar(Path(nested.name), budget, depth + 1)


def scan_image(image: str, saved_archive: Path, runner: Runner, timeout: int = 60) -> None:
    for command in (
        ["docker", "history", "--no-trunc", image],
        ["docker", "image", "inspect", image],
    ):
        _ensure_result_clean(runner.run(command, timeout=timeout))

    container = runner.run(["docker", "create", "--network", "none", image], timeout=timeout)
    container_id = container.stdout.strip()
    if not container_id:
        raise RuntimeError("docker create returned no container id")
    try:
        _ensure_result_clean(container)
        with tempfile.TemporaryDirectory() as directory:
            filesystem = Path(directory) / "filesystem.tar"
            _ensure_result_clean(
                runner.run(
                    ["docker", "export", "--output", str(filesystem), container_id],
                    timeout=timeout,
                )
            )
            _scan_tar(filesystem)
    finally:
        try:
            _ensure_result_clean(
                runner.run(["docker", "rm", "--force", container_id], timeout=timeout)
            )
        except BaseException as error:
            raise RuntimeError("canary scan container cleanup failed") from error
    _scan_tar(saved_archive)


def assert_clean_output(paths: Sequence[Path]) -> None:
    for path in paths:
        _scan_stream(path)


def _push_digest(output: str) -> str:
    digests = set(DIGEST_RE.findall(output))
    if len(digests) != 1:
        raise ValueError("registry push did not return exactly one valid digest")
    return digests.pop()


def _append(path: Path | None, text: str) -> None:
    if path is None:
        return
    with path.open("a", encoding="utf-8") as destination:
        destination.write(text)


def publish_image(
    *,
    archive: Path,
    checksum: Path,
    image_id: Path,
    local_image: str,
    commit: str,
    ref: str,
    actor: str,
    token: str,
    runner: Runner,
    output_path: Path | None = None,
    summary_path: Path | None = None,
) -> str:
    tags = release_tags(commit, ref)
    verify_checksum(archive, checksum)
    expected_image_id = read_image_id(image_id)
    loaded = runner.run(["docker", "load", "--input", str(archive)])
    load_lines = {line.strip() for line in (loaded.stdout + "\n" + loaded.stderr).splitlines()}
    if f"Loaded image: {local_image}" not in load_lines:
        raise ValueError("docker load did not report the expected local image")
    actual_image_id = inspect_image_id(local_image, runner)
    if actual_image_id != expected_image_id:
        raise ValueError("loaded image ID does not match tested image ID")
    revision = runner.run(
        [
            "docker",
            "image",
            "inspect",
            "--format",
            "{{ index .Config.Labels \"org.opencontainers.image.revision\" }}",
            local_image,
        ]
    ).stdout.strip()
    if revision != commit:
        raise ValueError("loaded image revision does not match source commit")

    if not actor or not token:
        raise ValueError("registry credentials are missing")
    for tag in tags:
        runner.run(["docker", "tag", local_image, tag])
    runner.run(
        ["docker", "login", "ghcr.io", "--username", actor, "--password-stdin"],
        input_text=token,
    )
    digests = []
    for tag in tags:
        result = runner.run(["docker", "push", tag])
        digests.append(_push_digest(result.stdout + "\n" + result.stderr))
    if len(set(digests)) != 1:
        raise ValueError("published tags returned different registry digests")

    digest = digests[0]
    reference = f"{IMAGE}@{digest}"
    _append(output_path, f"image={reference}\ndigest={digest}\nsource_commit={commit}\n")
    _append(
        summary_path,
        "## Published trex-core image\n\n"
        f"- Image: `{reference}`\n"
        f"- Digest: `{digest}`\n"
        f"- Source commit: `{commit}`\n",
    )
    return digest


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest="command", required=True)

    smoke = commands.add_parser("smoke")
    smoke.add_argument("--image", required=True)
    smoke.add_argument("--timeout", type=int, default=30)

    scan = commands.add_parser("scan")
    scan.add_argument("--image", required=True)
    scan.add_argument("--archive", type=Path, required=True)

    checksum = commands.add_parser("checksum")
    checksum.add_argument("--archive", type=Path, required=True)
    checksum.add_argument("--output", type=Path, required=True)

    record_id = commands.add_parser("record-image-id")
    record_id.add_argument("--image", required=True)
    record_id.add_argument("--output", type=Path, required=True)

    verify_id = commands.add_parser("verify-image-id")
    verify_id.add_argument("--image", required=True)
    verify_id.add_argument("--expected", type=Path, required=True)

    clean = commands.add_parser("assert-clean-output")
    clean.add_argument("paths", nargs="+", type=Path)

    publish = commands.add_parser("publish")
    publish.add_argument("--archive", type=Path, required=True)
    publish.add_argument("--checksum", type=Path, required=True)
    publish.add_argument("--image-id", type=Path, required=True)
    publish.add_argument("--local-image", required=True)
    publish.add_argument("--commit", required=True)
    publish.add_argument("--ref", required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    runner = SubprocessRunner()
    if args.command == "smoke":
        smoke_image(args.image, runner, args.timeout)
    elif args.command == "scan":
        scan_image(args.image, args.archive, runner)
    elif args.command == "checksum":
        write_checksum(args.archive, args.output)
    elif args.command == "record-image-id":
        write_image_id(args.image, args.output, runner)
    elif args.command == "verify-image-id":
        verify_image_id(args.image, args.expected, runner)
    elif args.command == "assert-clean-output":
        assert_clean_output(args.paths)
    elif args.command == "publish":
        publish_image(
            archive=args.archive,
            checksum=args.checksum,
            image_id=args.image_id,
            local_image=args.local_image,
            commit=args.commit,
            ref=args.ref,
            actor=os.environ.get("GITHUB_ACTOR", ""),
            token=os.environ.get("REGISTRY_TOKEN", ""),
            runner=runner,
            output_path=Path(path) if (path := os.environ.get("GITHUB_OUTPUT")) else None,
            summary_path=Path(path) if (path := os.environ.get("GITHUB_STEP_SUMMARY")) else None,
        )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError, RuntimeError, subprocess.SubprocessError) as error:
        print(f"image release failed: {error}", file=sys.stderr)
        raise SystemExit(1) from error