#!/usr/bin/env python3
"""Build, inspect, publish, verify and roll back trusted static releases."""

from __future__ import annotations

import argparse
import contextlib
import datetime as dt
import fcntl
import gzip
import hashlib
import io
import json
import os
from pathlib import Path, PurePosixPath
import re
import shutil
import sys
import tarfile
import tempfile
import urllib.error
import urllib.parse
import urllib.request


REVISION_RE = re.compile(r"^[0-9a-f]{40}$")
GENERATION_RE = re.compile(r"^[1-9][0-9]{0,11}$")
DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
MAX_FILES = 10_000
# The Sarah Weisman portfolio contains 164,651,446 bytes of site assets.
# Keep one bounded limit for the archive, receiver and extracted release.
MAX_BYTES = 200 * 1024 * 1024
SOURCE_MARKER = Path(".release-source.json")
PUBLIC_MARKER = Path(".well-known/release.json")
# Last desired-state selection this target applied, outside the served tree.
DESIRED_STATE = Path(".desired-state.json")
FORBIDDEN_PARTS = {".git", ".hg", ".svn"}


class ReleaseError(RuntimeError):
    pass


class SupersededError(ReleaseError):
    """A newer desired-state generation has already been applied here."""


def validate_revision(revision: str) -> None:
    if not REVISION_RE.fullmatch(revision):
        raise ReleaseError("revision must be a full 40-character lowercase Git SHA")


def validate_digest(digest: str) -> None:
    if not DIGEST_RE.fullmatch(digest):
        raise ReleaseError("digest must have the form sha256:<64 lowercase hex characters>")


def validate_order(order: dict) -> dict:
    """Validate the GitOps selection that requested a release.

    ``desired_generation`` is the ancestor count of ``desired_commit`` on the
    GitOps main branch. Main only moves to descendants, so a later selection
    always has a larger generation than an earlier one.
    """
    if not isinstance(order, dict) or set(order) != {"desired_commit", "desired_generation"}:
        raise ReleaseError("order must contain exactly desired_commit and desired_generation")
    if not isinstance(order["desired_commit"], str) or not REVISION_RE.fullmatch(order["desired_commit"]):
        raise ReleaseError("desired_commit must be a full 40-character lowercase Git SHA")
    if not isinstance(order["desired_generation"], str) or not GENERATION_RE.fullmatch(order["desired_generation"]):
        raise ReleaseError("desired_generation must be a positive decimal integer without leading zeros")
    return order


def load_desired_state(root: Path) -> dict | None:
    path = root / DESIRED_STATE
    if not path.exists() and not path.is_symlink():
        return None
    if path.is_symlink() or not path.is_file():
        raise ReleaseError("desired-state record must be a regular file")
    try:
        state = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ReleaseError(f"invalid desired-state record: {exc}") from exc
    if not isinstance(state, dict) or set(state) != {
        "desired_commit", "desired_generation", "revision", "artifact_digest",
    }:
        raise ReleaseError("desired-state record has an unsupported shape")
    validate_order({key: state[key] for key in ("desired_commit", "desired_generation")})
    validate_revision(state["revision"])
    validate_digest(state["artifact_digest"])
    return state


def check_order(root: Path, revision: str, digest: str, order: dict) -> None:
    """Refuse a request whose selection is older than the one applied here.

    An identical request (same generation, commit and release) is a retry of
    the accepted selection and may be applied again.
    """
    state = load_desired_state(root)
    if state is None:
        return
    requested = int(order["desired_generation"])
    applied = int(state["desired_generation"])
    if requested < applied:
        raise SupersededError(
            f"superseded: generation {requested} ({order['desired_commit']}) is older than "
            f"applied generation {applied} ({state['desired_commit']})"
        )
    if requested == applied and (
        order["desired_commit"] != state["desired_commit"]
        or revision != state["revision"]
        or digest != state["artifact_digest"]
    ):
        raise ReleaseError(
            f"generation {requested} was already applied with a different selection"
        )


def write_desired_state(root: Path, revision: str, digest: str, order: dict) -> None:
    state = {**order, "revision": revision, "artifact_digest": digest}
    temporary = root / f".desired-state.next-{os.getpid()}"
    temporary.write_text(json.dumps(state, sort_keys=True) + "\n", encoding="utf-8")
    os.replace(temporary, root / DESIRED_STATE)


def validate_site_path(path: PurePosixPath) -> None:
    if any(part in FORBIDDEN_PARTS for part in path.parts):
        raise ReleaseError(f"version-control metadata is forbidden: {path}")
    filename = path.name.lower()
    if filename == ".env" or filename.startswith(".env.") or filename in {"id_rsa", "id_ed25519"}:
        raise ReleaseError(f"common secret filename is forbidden: {path}")


def iter_files(root: Path):
    for path in sorted(root.rglob("*")):
        relative = PurePosixPath(path.relative_to(root).as_posix())
        validate_site_path(relative)
        if path.is_symlink():
            raise ReleaseError(f"symlinks are forbidden: {relative}")
        if path.is_file():
            yield path
        elif not path.is_dir():
            raise ReleaseError(f"unsupported filesystem entry: {relative}")


def validate_tree(root: Path) -> None:
    count = 0
    size = 0
    for path in iter_files(root):
        count += 1
        size += path.stat().st_size
        if count > MAX_FILES or size > MAX_BYTES:
            raise ReleaseError("site exceeds the file-count or uncompressed-size limit")
    if not (root / "index.html").is_file():
        raise ReleaseError("release has no index.html")


def load_source_marker(root: Path) -> dict:
    try:
        marker = json.loads((root / SOURCE_MARKER).read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ReleaseError(f"invalid or missing source marker: {exc}") from exc
    if set(marker) != {"revision"}:
        raise ReleaseError("source marker must contain only the revision")
    validate_revision(marker["revision"])
    validate_tree(root)
    return marker


def load_public_marker(root: Path) -> dict:
    try:
        marker = json.loads((root / PUBLIC_MARKER).read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ReleaseError(f"invalid or missing public release marker: {exc}") from exc
    if set(marker) != {"revision", "artifact_digest"}:
        raise ReleaseError("public marker must contain revision and artifact_digest")
    validate_revision(marker["revision"])
    validate_digest(marker["artifact_digest"])
    return marker


def build(source: Path, output: Path, revision: str) -> None:
    validate_revision(revision)
    if not source.is_dir():
        raise ReleaseError(f"source directory does not exist: {source}")
    if (source / SOURCE_MARKER).exists() or (source / PUBLIC_MARKER).exists():
        raise ReleaseError("source must not provide reserved release markers")
    validate_tree(source)

    source_root = source.resolve()
    output_root = output.resolve()
    if (
        output_root == source_root
        or output_root.is_relative_to(source_root)
        or source_root.is_relative_to(output_root)
    ):
        raise ReleaseError("build output and source must be separate directory trees")

    if output.exists():
        if output.is_symlink() or not output.is_dir():
            raise ReleaseError(f"build output must be a directory: {output}")
        shutil.rmtree(output)
    output.mkdir(parents=True)
    for path in iter_files(source):
        target = output / path.relative_to(source)
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(path, target)
        target.chmod(0o644)
    (output / SOURCE_MARKER).write_text(
        json.dumps({"revision": revision}, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    load_source_marker(output)


def package(directory: Path, archive: Path, checksum: Path) -> str:
    load_source_marker(directory)
    release_root = directory.resolve()
    if archive.resolve().is_relative_to(release_root) or checksum.resolve().is_relative_to(release_root):
        raise ReleaseError("archive and checksum must be written outside the release directory")
    if archive.resolve() == checksum.resolve():
        raise ReleaseError("archive and checksum paths must differ")
    archive.parent.mkdir(parents=True, exist_ok=True)
    with archive.open("wb") as raw:
        with gzip.GzipFile(filename="", mode="wb", fileobj=raw, mtime=0) as compressed:
            with tarfile.open(fileobj=compressed, mode="w", format=tarfile.PAX_FORMAT) as tar:
                for path in iter_files(directory):
                    relative = path.relative_to(directory).as_posix()
                    info = tar.gettarinfo(str(path), arcname=relative)
                    info.uid = info.gid = 0
                    info.uname = info.gname = ""
                    info.mode = 0o644
                    info.mtime = 0
                    with path.open("rb") as source:
                        tar.addfile(info, source)
    digest = "sha256:" + hashlib.sha256(archive.read_bytes()).hexdigest()
    checksum.write_text(f"{digest}  {archive.name}\n", encoding="ascii")
    return digest


def expected_archive_digest(checksum: Path, archive: Path) -> str:
    try:
        fields = checksum.read_text(encoding="ascii").strip().split()
    except OSError as exc:
        raise ReleaseError(f"cannot read checksum: {exc}") from exc
    if len(fields) != 2 or fields[1] != archive.name:
        raise ReleaseError("checksum file must contain 'sha256:<digest>  <archive-name>'")
    validate_digest(fields[0])
    return fields[0]


def inspect_archive(
    archive: Path,
    checksum: Path,
    revision: str,
) -> tuple[bytes, list[tarfile.TarInfo], str]:
    validate_revision(revision)
    expected = expected_archive_digest(checksum, archive)
    try:
        if archive.stat().st_size > MAX_BYTES:
            raise ReleaseError("compressed archive exceeds the size limit")
        archive_data = archive.read_bytes()
    except OSError as exc:
        raise ReleaseError(f"cannot read archive: {exc}") from exc
    actual = "sha256:" + hashlib.sha256(archive_data).hexdigest()
    if actual != expected:
        raise ReleaseError("archive digest mismatch")

    members: list[tarfile.TarInfo] = []
    member_names: set[str] = set()
    total = 0
    try:
        with tarfile.open(fileobj=io.BytesIO(archive_data), mode="r:gz") as tar:
            for member in tar.getmembers():
                path = PurePosixPath(member.name)
                if path.is_absolute() or not path.parts or any(part in ("", ".", "..") for part in path.parts):
                    raise ReleaseError(f"unsafe archive path: {member.name!r}")
                validate_site_path(path)
                if not member.isfile():
                    raise ReleaseError(f"archive may contain regular files only: {member.name!r}")
                if member.name in member_names:
                    raise ReleaseError(f"duplicate archive path: {member.name!r}")
                total += member.size
                members.append(member)
                member_names.add(member.name)
                if len(members) > MAX_FILES or total > MAX_BYTES:
                    raise ReleaseError("archive exceeds the file-count or uncompressed-size limit")
    except (OSError, tarfile.TarError) as exc:
        raise ReleaseError(f"invalid archive: {exc}") from exc

    names = {member.name for member in members}
    if "index.html" not in names or SOURCE_MARKER.as_posix() not in names:
        raise ReleaseError("archive must contain index.html and .release-source.json")
    if PUBLIC_MARKER.as_posix() in names:
        raise ReleaseError("artifact must not provide the publisher-owned public marker")
    with tempfile.TemporaryDirectory() as temp:
        destination = Path(temp)
        extract_members(archive_data, members, destination)
        marker = load_source_marker(destination)
        if marker["revision"] != revision:
            raise ReleaseError("archive revision does not match the requested revision")
    return archive_data, members, actual


def extract_members(archive_data: bytes, members: list[tarfile.TarInfo], destination: Path) -> None:
    with tarfile.open(fileobj=io.BytesIO(archive_data), mode="r:gz") as tar:
        by_name = {member.name: member for member in tar.getmembers()}
        for inspected in members:
            member = by_name.get(inspected.name)
            if member is None or not member.isfile() or member.size != inspected.size:
                raise ReleaseError("archive changed while it was being inspected")
            target = destination.joinpath(*PurePosixPath(member.name).parts)
            target.parent.mkdir(parents=True, exist_ok=True)
            source = tar.extractfile(member)
            if source is None:
                raise ReleaseError(f"cannot extract archive member: {member.name}")
            with source, target.open("wb") as output:
                shutil.copyfileobj(source, output)
            target.chmod(0o644)


def validate_verify_url(url: str) -> None:
    parsed = urllib.parse.urlsplit(url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc or parsed.query or parsed.fragment:
        raise ReleaseError("verify URL must be an HTTP(S) marker URL without query or fragment")
    if not parsed.path.endswith("/.well-known/release.json"):
        raise ReleaseError("verify URL must end with /.well-known/release.json")


def read_url_json(url: str, timeout: float) -> dict:
    request = urllib.request.Request(url, headers={"Cache-Control": "no-cache"})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            if response.status != 200:
                raise ReleaseError(f"public marker returned HTTP {response.status}")
            return json.loads(response.read().decode("utf-8"))
    except (OSError, urllib.error.URLError, json.JSONDecodeError) as exc:
        raise ReleaseError(f"public marker verification failed: {exc}") from exc


def verify_public(url: str, revision: str, digest: str, timeout: float = 10.0) -> dict:
    validate_revision(revision)
    validate_digest(digest)
    validate_verify_url(url)
    marker = read_url_json(url, timeout)
    if marker.get("revision") != revision or marker.get("artifact_digest") != digest:
        raise ReleaseError("public marker does not identify the expected release")

    index_url = url[: -len(PUBLIC_MARKER.as_posix())] + "index.html"
    try:
        with urllib.request.urlopen(
            urllib.request.Request(index_url, headers={"Cache-Control": "no-cache"}),
            timeout=timeout,
        ) as response:
            if response.status != 200 or not response.read(1):
                raise ReleaseError("public index is missing or empty")
    except (OSError, urllib.error.URLError) as exc:
        raise ReleaseError(f"public index verification failed: {exc}") from exc
    return marker


@contextlib.contextmanager
def release_lock(root: Path):
    if (
        not root.is_absolute()
        or root == Path("/")
        or root.resolve(strict=False) != root
        or not root.is_dir()
        or root.is_symlink()
    ):
        raise ReleaseError(f"release root must be a provisioned directory: {root}")
    with (root / ".deploy.lock").open("a+") as handle:
        fcntl.flock(handle, fcntl.LOCK_EX)
        yield


def current_release(root: Path) -> dict | None:
    pointer = root / "current"
    if not pointer.exists() and not pointer.is_symlink():
        return None
    if not pointer.is_symlink():
        raise ReleaseError(f"current target must be a managed symlink: {pointer}")
    resolved = pointer.resolve(strict=True)
    releases = (root / "releases").resolve()
    try:
        relative = resolved.relative_to(releases)
    except ValueError as exc:
        raise ReleaseError("current pointer escapes the managed releases directory") from exc
    if len(relative.parts) != 1 or not re.fullmatch(r"[0-9a-f]{64}", relative.name):
        raise ReleaseError("current pointer has an invalid release target")
    marker = load_public_marker(resolved)
    if marker["artifact_digest"] != f"sha256:{relative.name}":
        raise ReleaseError("current pointer and public marker digest disagree")
    return marker


def switch(root: Path, digest: str) -> None:
    digest_hex = digest.removeprefix("sha256:")
    temporary = root / f".current.next-{os.getpid()}"
    if temporary.exists() or temporary.is_symlink():
        temporary.unlink()
    temporary.symlink_to(Path("releases") / digest_hex)
    os.replace(temporary, root / "current")


def record(
    root: Path,
    action: str,
    revision: str,
    digest: str,
    previous: dict | None,
    outcome: str,
    actor: str | None,
    order: dict | None = None,
) -> None:
    entry = {
        "at": dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z"),
        "action": action,
        "actor": actor or os.environ.get("GITHUB_ACTOR") or os.environ.get("USER") or "unknown",
        "artifact_digest": digest,
        "revision": revision,
        "previous_revision": previous["revision"] if previous else None,
        "previous_artifact_digest": previous["artifact_digest"] if previous else None,
        "target": str(root / "current"),
        "outcome": outcome,
        "desired_commit": order["desired_commit"] if order else None,
        "desired_generation": order["desired_generation"] if order else None,
    }
    with (root / ".deployments.jsonl").open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(entry, sort_keys=True) + "\n")


def promote(
    root: Path,
    revision: str,
    digest: str,
    *,
    verify_url: str | None,
    action: str,
    actor: str | None,
    order: dict | None = None,
) -> None:
    validate_revision(revision)
    validate_digest(digest)
    release = root / "releases" / digest.removeprefix("sha256:")
    if not release.is_dir():
        raise ReleaseError(f"release is not staged: {digest}")
    source = load_source_marker(release)
    marker = load_public_marker(release)
    if source["revision"] != revision or marker != {"revision": revision, "artifact_digest": digest}:
        raise ReleaseError("staged release identity mismatch")

    previous = current_release(root)
    switch(root, digest)
    try:
        if verify_url:
            verify_public(verify_url, revision, digest)
        record(root, action, revision, digest, previous, "verified", actor, order)
    except BaseException:
        if previous is None:
            (root / "current").unlink(missing_ok=True)
        else:
            switch(root, previous["artifact_digest"])
        record(root, action, revision, digest, previous, "rolled-back", actor, order)
        raise
    if order is not None:
        write_desired_state(root, revision, digest, order)


def publish(
    archive: Path,
    checksum: Path,
    root: Path,
    revision: str,
    verify_url: str | None,
    actor: str | None = None,
    order: dict | None = None,
) -> None:
    if order is not None:
        validate_order(order)
    archive_data, members, digest = inspect_archive(archive, checksum, revision)
    with release_lock(root):
        if order is not None:
            try:
                check_order(root, revision, digest, order)
            except SupersededError:
                record(root, "publish", revision, digest, current_release(root), "superseded", actor, order)
                raise
        releases = root / "releases"
        releases.mkdir(exist_ok=True)
        destination = releases / digest.removeprefix("sha256:")
        if destination.exists():
            marker = load_public_marker(destination)
            if marker != {"revision": revision, "artifact_digest": digest}:
                raise ReleaseError("existing immutable release has the wrong identity")
        else:
            with tempfile.TemporaryDirectory(prefix=".incoming-", dir=releases) as temp:
                incoming = Path(temp)
                extract_members(archive_data, members, incoming)
                source = load_source_marker(incoming)
                if source["revision"] != revision:
                    raise ReleaseError("extracted release revision mismatch")
                marker_path = incoming / PUBLIC_MARKER
                marker_path.parent.mkdir(parents=True, exist_ok=True)
                marker_path.write_text(
                    json.dumps({"revision": revision, "artifact_digest": digest}, separators=(",", ":")) + "\n",
                    encoding="utf-8",
                )
                load_public_marker(incoming)
                os.replace(incoming, destination)
        promote(
            root,
            revision,
            digest,
            verify_url=verify_url,
            action="publish",
            actor=actor,
            order=order,
        )


def rollback(
    root: Path,
    revision: str,
    digest: str,
    verify_url: str | None,
    actor: str | None = None,
) -> None:
    with release_lock(root):
        promote(root, revision, digest, verify_url=verify_url, action="rollback", actor=actor)


def retention(root: Path, keep: int) -> dict:
    """Report releases outside the retention set. Never deletes anything.

    Kept: the current release, the last applied desired selection, and the
    ``keep`` most recent distinct verified releases from the deployment log.
    """
    if keep < 1:
        raise ReleaseError("keep must be at least 1")
    releases = root / "releases"
    present = sorted(
        path.name for path in releases.iterdir()
        if path.is_dir() and re.fullmatch(r"[0-9a-f]{64}", path.name)
    ) if releases.is_dir() else []
    reasons: dict[str, list[str]] = {}

    def keep_digest(digest: str | None, reason: str) -> None:
        if digest:
            reasons.setdefault(digest.removeprefix("sha256:"), []).append(reason)

    current = current_release(root)
    keep_digest(current and current["artifact_digest"], "current")
    state = load_desired_state(root)
    keep_digest(state and state["artifact_digest"], "desired-state")
    verified: list[str] = []
    log = root / ".deployments.jsonl"
    if log.is_file():
        for line in reversed(log.read_text(encoding="utf-8").splitlines()):
            try:
                entry = json.loads(line)
            except json.JSONDecodeError:
                continue
            digest = entry.get("artifact_digest")
            if entry.get("outcome") == "verified" and digest not in verified:
                verified.append(digest)
    for digest in verified[:keep]:
        keep_digest(digest, "recent-verified")
    return {
        "root": str(root),
        "keep": {digest: reasons[digest] for digest in present if digest in reasons},
        "candidates": [digest for digest in present if digest not in reasons],
    }


def parser() -> argparse.ArgumentParser:
    command = argparse.ArgumentParser(description=__doc__)
    sub = command.add_subparsers(dest="command", required=True)

    build_cmd = sub.add_parser("build")
    build_cmd.add_argument("--source", type=Path, required=True)
    build_cmd.add_argument("--output", type=Path, required=True)
    build_cmd.add_argument("--revision", required=True)

    package_cmd = sub.add_parser("package")
    package_cmd.add_argument("--directory", type=Path, required=True)
    package_cmd.add_argument("--archive", type=Path, required=True)
    package_cmd.add_argument("--checksum", type=Path, required=True)

    inspect_cmd = sub.add_parser("inspect")
    inspect_cmd.add_argument("--archive", type=Path, required=True)
    inspect_cmd.add_argument("--checksum", type=Path, required=True)
    inspect_cmd.add_argument("--revision", required=True)

    publish_cmd = sub.add_parser("publish")
    publish_cmd.add_argument("--archive", type=Path, required=True)
    publish_cmd.add_argument("--checksum", type=Path, required=True)
    publish_cmd.add_argument("--root", type=Path, required=True)
    publish_cmd.add_argument("--revision", required=True)
    publish_cmd.add_argument("--verify-url")
    publish_cmd.add_argument("--actor")

    verify_cmd = sub.add_parser("verify")
    verify_cmd.add_argument("--url", required=True)
    verify_cmd.add_argument("--revision", required=True)
    verify_cmd.add_argument("--digest", required=True)

    rollback_cmd = sub.add_parser("rollback")
    rollback_cmd.add_argument("--root", type=Path, required=True)
    rollback_cmd.add_argument("--revision", required=True)
    rollback_cmd.add_argument("--digest", required=True)
    rollback_cmd.add_argument("--verify-url")
    rollback_cmd.add_argument("--actor")

    retention_cmd = sub.add_parser("retention", help="report prune candidates; never deletes")
    retention_cmd.add_argument("--root", type=Path, required=True)
    retention_cmd.add_argument("--keep", type=int, default=5)
    return command


def main() -> int:
    args = parser().parse_args()
    try:
        if args.command == "build":
            build(args.source, args.output, args.revision)
        elif args.command == "package":
            package(args.directory, args.archive, args.checksum)
        elif args.command == "inspect":
            inspect_archive(args.archive, args.checksum, args.revision)
        elif args.command == "publish":
            publish(args.archive, args.checksum, args.root, args.revision, args.verify_url, args.actor)
        elif args.command == "verify":
            verify_public(args.url, args.revision, args.digest)
        elif args.command == "rollback":
            rollback(args.root, args.revision, args.digest, args.verify_url, args.actor)
        elif args.command == "retention":
            print(json.dumps(retention(args.root, args.keep), indent=2, sort_keys=True))
    except ReleaseError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())