#!/usr/bin/env python3
"""Fail-closed PortfolioHub image verification and publication helpers."""
from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import re
import subprocess
import sys
import tarfile
import tempfile
import time
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

IMAGE = "ghcr.io/fourblocks2020/portfolio-hub"
CANARY = b"PORTFOLIO_HUB_SECRET_CANARY_DO_NOT_SHIP"
COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
VERSION_REF_RE = re.compile(r"^refs/tags/v(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)$")
IMAGE_ID_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
DIGEST_RE = re.compile(r"(?:^|\s)digest:\s*(sha256:[0-9a-f]{64})(?:\s|$)")

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

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:
        result = subprocess.run(command, input=input_text, text=True, capture_output=True,
                                check=True, timeout=timeout)
        return Result(result.stdout, result.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 validate_checkout_identity(actual: str, expected: str) -> str:
    validate_commit(actual); validate_commit(expected)
    if actual != expected:
        raise ValueError("checked-out HEAD does not match the tagged source commit")
    return actual

def validate_release_event(event_name: str, ref: str, commit: str) -> str:
    validate_commit(commit)
    if event_name != "push" or not VERSION_REF_RE.fullmatch(ref):
        raise ValueError("publication requires a push of a strict vX.Y.Z tag")
    return ref.removeprefix("refs/tags/")

def release_tags(event_name: str, ref: str, commit: str) -> tuple[str, str]:
    version = validate_release_event(event_name, ref, commit)
    return f"{IMAGE}:{commit[:12]}", f"{IMAGE}:{version}"

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, output: Path) -> str:
    checksum = sha256_file(archive)
    output.write_text(f"{checksum}  {archive.name}\n", encoding="ascii")
    return checksum

def verify_checksum(archive: Path, checksum_path: Path) -> str:
    if not archive.is_file() or not checksum_path.is_file():
        raise ValueError("image archive or checksum evidence is missing")
    fields = checksum_path.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")
    return fields[0]

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

def inspect_value(image: str, template: str, runner: Runner) -> str:
    return _clean(runner.run(["docker", "image", "inspect", "--format", template, image])).stdout.strip()

def inspect_image_id(image: str, runner: Runner) -> str:
    value = inspect_value(image, "{{.Id}}", runner)
    if not IMAGE_ID_RE.fullmatch(value):
        raise ValueError("image ID has an invalid format")
    return value

def inspect_revision(image: str, runner: Runner) -> str:
    return inspect_value(image, '{{ index .Config.Labels "org.opencontainers.image.revision" }}', runner)

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

def read_image_id(path: Path) -> str:
    value = path.read_text(encoding="ascii").strip()
    if not IMAGE_ID_RE.fullmatch(value):
        raise ValueError("tested image ID evidence has an invalid format")
    return value

def _error_text(error: subprocess.CalledProcessError) -> str:
    return "\n".join(str(value or "") for value in (error.stdout, error.stderr)).lower()

def _remove(command: list[str], absent: tuple[str, ...], runner: Runner, timeout: int) -> None:
    try:
        _clean(runner.run(command, timeout=timeout))
    except subprocess.CalledProcessError as error:
        if not any(marker in _error_text(error) for marker in absent):
            raise

def _cleanup(names: list[str], network: str | None, volume: str | None,
             runner: Runner, timeout: int) -> None:
    errors: list[BaseException] = []
    for name in reversed(names):
        try:
            _remove(["docker", "rm", "--force", name], ("no such container",), runner, timeout)
        except BaseException as error: errors.append(error)
    if network:
        try:
            _remove(["docker", "network", "rm", network], ("not found", "no such network"), runner, timeout)
        except BaseException as error: errors.append(error)
    if volume:
        try:
            _remove(["docker", "volume", "rm", "--force", volume], ("no such volume", "not found"), runner, timeout)
        except BaseException as error: errors.append(error)
    if errors:
        raise RuntimeError("image smoke cleanup failed") from errors[0]

def smoke_installed_image(image: str, expected_inventory: tuple[str, ...], runner: Runner,
                          timeout: int = 30) -> None:
    if not expected_inventory:
        raise ValueError("expected migration inventory must be non-empty")
    inventory_code = (
        "import json; from app.db.migrations import migration_inventory; "
        f"expected={json.dumps(list(expected_inventory))}; actual=list(migration_inventory()); "
        "assert actual == expected, (actual, expected)"
    )
    probes = (("python", "-I", "-c", "import app"), ("pfhub", "--help"),
              ("python", "-I", "-c", inventory_code))
    for index, probe in enumerate(probes):
        name = f"portfolio-hub-install-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: _clean(runner.run(command, timeout=timeout))
        except BaseException as error: failure = error
        try: _cleanup([name], None, None, runner, timeout)
        except BaseException as error:
            if failure is None: failure = error
        if failure: raise RuntimeError(f"install smoke probe {index + 1} failed") from failure

def _wait(command: list[str], runner: Runner, deadline: float, description: str) -> Result:
    last: BaseException | None = None
    while time.monotonic() < deadline:
        try: return _clean(runner.run(command, timeout=5))
        except (OSError, subprocess.SubprocessError) as error:
            last = error; time.sleep(1)
    raise RuntimeError(f"timed out waiting for {description}") from last

def startup_smoke(image: str, runner: Runner, timeout: int = 120) -> dict:
    """Run DB-disabled health, then isolated PostgreSQL migration, ping, and health."""
    suffix = f"{os.getpid()}-{time.time_ns()}"
    plain = f"portfolio-hub-plain-{suffix}"; db = f"portfolio-hub-db-{suffix}"
    migrate = f"portfolio-hub-migrate-{suffix}"; app = f"portfolio-hub-app-{suffix}"
    network = f"portfolio-hub-smoke-{suffix}"; volume = f"portfolio-hub-db-{suffix}"
    names: list[str] = []
    deadline = time.monotonic() + timeout
    failure: BaseException | None = None; result: dict = {}
    db_env = ["-e", "PORTFOLIOHUB_DB_ENABLED=true", "-e", "TIMESCALEDB_HOST=postgres",
              "-e", "TIMESCALEDB_PORT=5432", "-e", "TIMESCALEDB_NAME=smoke",
              "-e", "TIMESCALEDB_USER=smoke", "-e", "TIMESCALEDB_PASSWORD=synthetic-db-only",
              "-e", "TIMESCALEDB_SSLMODE=disable", "-e", "CREDENTIAL_ENCRYPTION_KEY=bnVsbC1vbmx5LXN5bnRoZXRpYy1rZXktMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwPQ=="]
    http_probe = "import json,urllib.request; d=json.load(urllib.request.urlopen('http://TARGET:8000/health')); assert d['status']=='ok'; print(json.dumps(d,sort_keys=True))"
    try:
        runner.run(["docker", "network", "create", "--internal", network], timeout=10)
        runner.run(["docker", "volume", "create", volume], timeout=10)
        names.append(plain)
        runner.run(["docker", "run", "-d", "--name", plain, "--network", network,
                    "--network-alias", "plain", "--read-only", "--cap-drop", "ALL",
                    "--security-opt", "no-new-privileges", "--tmpfs", "/tmp:rw,noexec,nosuid,size=16m",
                    "-e", "PORTFOLIOHUB_DB_ENABLED=false", image], timeout=20)
        plain_result = _wait(["docker", "run", "--rm", "--network", network, "--entrypoint", "python",
                              image, "-c", http_probe.replace("TARGET", "plain")], runner, deadline, "DB-disabled /health")
        plain_payload = json.loads(plain_result.stdout)
        if not isinstance(plain_payload, dict) or plain_payload.get("db_enabled") is not False:
            raise RuntimeError("DB-disabled health payload has the wrong shape")
        names.append(db)
        runner.run(["docker", "run", "-d", "--name", db, "--network", network,
                    "--network-alias", "postgres", "-e", "POSTGRES_DB=smoke", "-e", "POSTGRES_USER=smoke",
                    "-e", "POSTGRES_PASSWORD=synthetic-db-only", "--mount", f"type=volume,src={volume},dst=/var/lib/postgresql/data",
                    "postgres:16-alpine"], timeout=30)
        _wait(["docker", "exec", db, "pg_isready", "-U", "smoke", "-d", "smoke"], runner, deadline, "PostgreSQL")
        names.append(migrate)
        runner.run(["docker", "run", "--name", migrate, "--network", network, *db_env,
                    "--entrypoint", "pfhub", image, "migrate-db"], timeout=60)
        check = "import psycopg; c=psycopg.connect(host='postgres',dbname='smoke',user='smoke',password='synthetic-db-only'); r=c.execute(\"select version_num from alembic_version\").fetchone(); assert r and r[0]; assert c.execute('select 1').fetchone()==(1,); print(r[0])"
        head = _clean(runner.run(["docker", "run", "--rm", "--network", network, "--entrypoint", "python",
                                  image, "-c", check], timeout=30)).stdout.strip()
        if not head: raise RuntimeError("Alembic head/DB ping returned no revision")
        names.append(app)
        runner.run(["docker", "run", "-d", "--name", app, "--network", network, "--network-alias", "app",
                    "--read-only", "--cap-drop", "ALL", "--security-opt", "no-new-privileges",
                    "--tmpfs", "/tmp:rw,noexec,nosuid,size=16m", *db_env, image], timeout=20)
        app_result = _wait(["docker", "run", "--rm", "--network", network, "--entrypoint", "python",
                            image, "-c", http_probe.replace("TARGET", "app")], runner, deadline, "DB-enabled /health")
        app_payload = json.loads(app_result.stdout)
        if not isinstance(app_payload, dict) or app_payload.get("db_enabled") is not True:
            raise RuntimeError("DB-enabled health payload has the wrong shape")
        user = inspect_value(image, "{{.Config.User}}", runner)
        if user not in {"10001", "10001:10001"}: raise RuntimeError("image default user must be 10001")
        result = {"plain_health": "ok", "db_health": "ok", "db_ping": "ok", "alembic_head": head}
    except BaseException as error: failure = error
    try: _cleanup(names, network, volume, runner, 15)
    except BaseException as error:
        if failure is None: failure = error
    if failure:
        if isinstance(failure, (KeyboardInterrupt, SystemExit)): raise failure
        raise RuntimeError("isolated PortfolioHub startup smoke failed") from failure
    return result

def _scan_bytes(data: bytes) -> None:
    if CANARY in data: raise ValueError("canary marker found on a scanned surface")

def _scan_tar(data: bytes, nested: bool) -> None:
    _scan_bytes(data)
    try:
        with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as archive:
            for member in archive:
                metadata = [member.name, member.linkname, member.uname, member.gname]
                for key, value in member.pax_headers.items(): metadata.extend((key, value))
                _scan_bytes("\n".join(metadata).encode(errors="replace"))
                if member.isfile() and (source := archive.extractfile(member)):
                    payload = source.read(); _scan_bytes(payload)
                    if nested: _scan_tar(payload, False)
    except (tarfile.TarError, OSError, EOFError): pass

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

    # Docker can create the container successfully and still return output that
    # fails the canary scan. Retain its identity before scanning the complete
    # result so every post-create failure is cleanup-protected.
    created = runner.run(["docker", "create", "--network", "none", image], timeout=timeout)
    container = created.stdout.strip().splitlines()[0] if created.stdout.strip() else ""
    if not container:
        _clean(created)
        raise RuntimeError("docker create returned no container ID")

    failure: BaseException | None = None
    try:
        _clean(created)
        with tempfile.TemporaryDirectory() as directory:
            filesystem = Path(directory) / "filesystem.tar"
            _clean(runner.run(["docker", "export", "--output", str(filesystem), container], timeout=timeout))
            _scan_tar(filesystem.read_bytes(), False)
    except BaseException as error:
        failure = error
    try:
        _cleanup([container], None, None, runner, timeout)
    except BaseException as error:
        if failure is None:
            failure = error
    if failure:
        raise failure
    _scan_tar(saved_archive.read_bytes(), True)

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

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:
        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, event_name: str, ref: str, actor: str, token: str,
                  runner: Runner, output_path: Path | None = None,
                  summary_path: Path | None = None) -> str:
    tags = release_tags(event_name, ref, commit)
    archive_sha = verify_checksum(archive, checksum); expected_id = read_image_id(image_id)
    loaded = _clean(runner.run(["docker", "load", "--input", str(archive)]))
    if f"Loaded image: {local_image}" not in (loaded.stdout + loaded.stderr).splitlines():
        raise ValueError("docker load did not report the expected local image")
    actual_id = inspect_image_id(local_image, runner)
    if actual_id != expected_id: raise ValueError("loaded image ID does not match tested image ID")
    if inspect_revision(local_image, runner) != commit: raise ValueError("loaded image revision does not match tagged 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:
        pushed = _clean(runner.run(["docker", "push", tag])); digests.append(_push_digest(pushed.stdout + "\n" + pushed.stderr))
    if len(set(digests)) != 1: raise ValueError("published selectors returned different registry digests")
    digest = digests[0]; immutable = f"{IMAGE}@{digest}"
    evidence = (f"image={immutable}\ndigest={digest}\nsource_commit={commit}\nimage_id={actual_id}\n"
                f"archive_sha256={archive_sha}\nsource_selector={tags[0]}\nversion_selector={tags[1]}\n")
    _append(output_path, evidence)
    _append(summary_path, "## Published PortfolioHub image\n\n" + "".join(f"- {line}\n" for line in evidence.splitlines()))
    return digest

def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(); commands = parser.add_subparsers(dest="command", required=True)
    event = commands.add_parser("validate-event"); event.add_argument("--event", required=True); event.add_argument("--ref", required=True); event.add_argument("--commit", required=True)
    checkout = commands.add_parser("validate-checkout"); checkout.add_argument("--actual", required=True); checkout.add_argument("--expected", required=True)
    smoke = commands.add_parser("smoke"); smoke.add_argument("--image", required=True); smoke.add_argument("--inventory", type=Path, required=True); smoke.add_argument("--timeout", type=int, default=120)
    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 = commands.add_parser("record-image-id"); record.add_argument("--image", required=True); record.add_argument("--output", type=Path, required=True)
    clean = commands.add_parser("assert-clean-output"); clean.add_argument("paths", nargs="+", type=Path)
    publish = commands.add_parser("publish")
    for option in ("archive", "checksum", "image-id"): publish.add_argument(f"--{option}", type=Path, required=True)
    for option in ("local-image", "commit", "event", "ref"): publish.add_argument(f"--{option}", required=True)
    return parser

def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv); runner = SubprocessRunner()
    if args.command == "validate-event": validate_release_event(args.event, args.ref, args.commit)
    elif args.command == "validate-checkout": validate_checkout_identity(args.actual, args.expected)
    elif args.command == "smoke":
        inventory = tuple(line for line in args.inventory.read_text().splitlines() if line)
        smoke_installed_image(args.image, inventory, runner, args.timeout); startup_smoke(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 == "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, event_name=args.event,
                      ref=args.ref, actor=os.environ.get("GITHUB_ACTOR", ""),
                      token=os.environ.get("REGISTRY_TOKEN", ""), runner=runner,
                      output_path=Path(value) if (value := os.environ.get("GITHUB_OUTPUT")) else None,
                      summary_path=Path(value) if (value := 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