#!/usr/bin/env python3
"""Promote and roll back immutable Timmy staging release archives safely."""
from __future__ import annotations

import argparse
from contextlib import contextmanager
from dataclasses import dataclass
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import re
import secrets
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
from typing import Callable, Sequence
import urllib.error
import urllib.request

COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
TAG_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,79}$")
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
FORBIDDEN_SUFFIXES = (".pem", ".key", ".p12", ".pfx", ".gguf", ".bin", ".safetensors", ".onnx", ".pyc")
FORBIDDEN_COMPONENTS = {".env", ".git", ".ssh", "secrets", "credentials", "__pycache__"}
FORBIDDEN_PREFIXES = (("video",), ("artifacts",), ("research", "source-pages"))


class DeploymentError(RuntimeError):
    """A bounded, operator-safe deployment failure."""


@dataclass(frozen=True)
class DeploymentConfig:
    root: Path
    restart_command: tuple[str, ...]
    stop_command: tuple[str, ...]
    smoke_command: tuple[str, ...]
    health_url: str
    command_timeout: float = 180.0
    health_timeout: float = 30.0
    max_members: int = 10_000
    max_member_bytes: int = 50 * 1024 * 1024
    max_total_bytes: int = 250 * 1024 * 1024

    @property
    def releases(self) -> Path:
        return self.root / "releases"

    @property
    def current(self) -> Path:
        return self.root / "current"


RunCommand = Callable[..., subprocess.CompletedProcess]
HealthCheck = Callable[[str, str, float], dict]


def _validate_identity(tag: str, commit: str, expected_sha256: str) -> None:
    if not TAG_RE.fullmatch(tag):
        raise DeploymentError("tag must be 1-80 safe release characters")
    if not COMMIT_RE.fullmatch(commit):
        raise DeploymentError("commit must be exactly 40 lowercase hexadecimal characters")
    if not SHA256_RE.fullmatch(expected_sha256):
        raise DeploymentError("SHA-256 must be exactly 64 lowercase hexadecimal characters")


@contextmanager
def _verified_archive_copy(source_path: Path, expected_sha256: str):
    """Copy and hash an untrusted regular archive once, then yield the private copy."""
    source_fd = -1
    flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)
    try:
        source_fd = os.open(source_path, flags)
        if not stat.S_ISREG(os.fstat(source_fd).st_mode):
            raise DeploymentError("archive source must be a regular file, not a symlink or special file")
        with tempfile.TemporaryDirectory(prefix="timmy-verified-archive-") as private_dir:
            private_path = Path(private_dir) / "archive"
            digest = hashlib.sha256()
            with os.fdopen(source_fd, "rb", closefd=True) as source:
                source_fd = -1
                with private_path.open("xb") as destination:
                    for chunk in iter(lambda: source.read(1024 * 1024), b""):
                        digest.update(chunk)
                        destination.write(chunk)
                    destination.flush()
                    os.fsync(destination.fileno())
            private_path.chmod(0o400)
            if digest.hexdigest() != expected_sha256:
                raise DeploymentError("archive SHA-256 mismatch; refusing extraction")
            yield private_path
    except DeploymentError:
        raise
    except OSError as error:
        raise DeploymentError(f"cannot stage archive: {error.strerror or 'I/O error'}") from error
    finally:
        if source_fd >= 0:
            os.close(source_fd)


def _safe_parts(name: str) -> tuple[str, ...]:
    if not name or "\\" in name or "\x00" in name or name.startswith("/") or "//" in name:
        raise DeploymentError(f"unsafe archive path: {name!r}")
    path = PurePosixPath(name)
    if path.is_absolute() or any(part in ("", ".", "..") for part in path.parts):
        raise DeploymentError(f"unsafe archive path: {name!r}")
    return path.parts


def _is_forbidden(parts: tuple[str, ...]) -> bool:
    lowered = tuple(part.lower() for part in parts)
    basename = lowered[-1]
    return (
        any(part in FORBIDDEN_COMPONENTS or part.startswith(".env.") for part in lowered)
        or basename in {"id_rsa", "id_ed25519", "authorized_keys"}
        or basename.endswith(FORBIDDEN_SUFFIXES)
        or any(lowered[:len(prefix)] == prefix for prefix in FORBIDDEN_PREFIXES)
    )


def inspect_archive(archive_path: Path, config: DeploymentConfig) -> tuple[list[tarfile.TarInfo], str | None]:
    """Validate every tar entry and return members plus a common wrapper directory."""
    try:
        with tarfile.open(archive_path, mode="r:*") as archive:
            members = archive.getmembers()
    except (OSError, tarfile.TarError) as error:
        raise DeploymentError("archive is not a readable tar file") from error
    if not members:
        raise DeploymentError("archive is empty")
    if len(members) > config.max_members:
        raise DeploymentError("archive member count exceeds configured limit")
    total = 0
    all_parts: list[tuple[str, ...]] = []
    for member in members:
        parts = _safe_parts(member.name)
        all_parts.append(parts)
        if not (member.isdir() or member.isreg()):
            raise DeploymentError(f"unsupported archive member type: {member.name}")
        if member.size < 0 or member.size > config.max_member_bytes:
            raise DeploymentError(f"archive member size exceeds configured limit: {member.name}")
        total += member.size
        if total > config.max_total_bytes:
            raise DeploymentError("archive expanded size exceeds configured limit")
    common_root = all_parts[0][0] if all(parts[0] == all_parts[0][0] for parts in all_parts) else None
    if common_root and _is_forbidden((common_root,)):
        raise DeploymentError(f"forbidden archive artifact: {common_root}")
    for parts in all_parts:
        relative = parts[1:] if common_root else parts
        if not relative:
            continue
        if _is_forbidden(relative):
            raise DeploymentError(f"forbidden archive artifact: {'/'.join(relative)}")
    return members, common_root


def _extract_validated(archive_path: Path, destination: Path, members: list[tarfile.TarInfo], common_root: str | None) -> None:
    with tarfile.open(archive_path, mode="r:*") as archive:
        for member in members:
            parts = _safe_parts(member.name)
            relative = parts[1:] if common_root else parts
            if not relative:
                continue
            target = destination.joinpath(*relative)
            if member.isdir():
                target.mkdir(parents=True, exist_ok=True, mode=0o755)
                continue
            target.parent.mkdir(parents=True, exist_ok=True, mode=0o755)
            source = archive.extractfile(member)
            if source is None:
                raise DeploymentError(f"could not read archive member: {member.name}")
            with source, target.open("xb") as output:
                shutil.copyfileobj(source, output, length=1024 * 1024)
            target.chmod(0o755 if member.mode & 0o111 else 0o644)


def _ensure_deployment_root(config: DeploymentConfig, *, create: bool) -> None:
    if config.root.is_symlink():
        raise DeploymentError("deployment root must not be a symlink")
    if config.root.exists() and not config.root.is_dir():
        raise DeploymentError("deployment root is not a directory")
    if create:
        config.root.mkdir(parents=True, exist_ok=True, mode=0o755)
        if config.root.is_symlink() or not config.root.is_dir():
            raise DeploymentError("deployment root is not a safe directory")


def _ensure_releases_directory(config: DeploymentConfig, *, create: bool) -> None:
    _ensure_deployment_root(config, create=create)
    if config.releases.is_symlink():
        raise DeploymentError("releases directory must not be a symlink")
    if config.releases.exists() and not config.releases.is_dir():
        raise DeploymentError("releases directory is not a directory")
    if create:
        config.releases.mkdir(parents=True, exist_ok=True, mode=0o755)


def _current_commit(config: DeploymentConfig) -> str | None:
    _ensure_releases_directory(config, create=False)
    if not config.current.is_symlink():
        if config.current.exists():
            raise DeploymentError("current exists but is not a symlink")
        return None
    target = os.readlink(config.current)
    target_path = Path(target)
    resolved = (config.current.parent / target_path).resolve() if not target_path.is_absolute() else target_path.resolve()
    try:
        relative = resolved.relative_to(config.releases.resolve())
    except ValueError as error:
        raise DeploymentError("current symlink escapes the releases directory") from error
    if len(relative.parts) != 1 or not COMMIT_RE.fullmatch(relative.name) or not resolved.is_dir():
        raise DeploymentError("current symlink does not name a valid immutable release")
    return relative.name


def _atomic_point(config: DeploymentConfig, commit: str | None) -> None:
    _ensure_releases_directory(config, create=False)
    if config.current.exists() and not config.current.is_symlink():
        raise DeploymentError("current exists but is not a symlink")
    temporary = config.root / f".current-{os.getpid()}-{secrets.token_hex(6)}"
    try:
        if commit is None:
            config.current.unlink(missing_ok=True)
            return
        temporary.symlink_to(Path("releases") / commit)
        os.replace(temporary, config.current)
    finally:
        temporary.unlink(missing_ok=True)


def _run_argv(argv: Sequence[str], **options) -> subprocess.CompletedProcess:
    if not argv:
        raise DeploymentError("configured command must be a non-empty argv array")
    timeout = float(options.get("timeout", 180.0))
    return subprocess.run(list(argv), check=True, text=True, capture_output=True, shell=False, timeout=timeout)


def poll_health(url: str, expected_commit: str, timeout: float) -> dict:
    deadline = time.monotonic() + timeout
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise DeploymentError("health check deadline expired")
        try:
            request = urllib.request.Request(url, headers={"accept": "application/json"})
            with urllib.request.urlopen(request, timeout=min(2.0, remaining)) as response:
                raw = response.read(65_537)
                if len(raw) > 65_536:
                    raise DeploymentError("health response exceeded 64 KiB")
                payload = json.loads(raw)
                if response.status == 200 and payload.get("ok") is True and payload.get("commit") == expected_commit:
                    return payload
        except DeploymentError:
            raise
        except (OSError, urllib.error.URLError, json.JSONDecodeError, AttributeError):
            pass
        time.sleep(min(0.2, max(0.0, deadline - time.monotonic())))


def _verify(config: DeploymentConfig, commit: str, run_command: RunCommand, health_check: HealthCheck, *, smoke: bool) -> None:
    run_command(config.restart_command, check=True, text=True, capture_output=True, shell=False, timeout=config.command_timeout)
    health_check(config.health_url, commit, config.health_timeout)
    if smoke:
        run_command(config.smoke_command, check=True, text=True, capture_output=True, shell=False, timeout=config.command_timeout)


def promote(*, config: DeploymentConfig, tag: str, archive: Path, expected_sha256: str, commit: str,
            run_command: RunCommand = _run_argv, health_check: HealthCheck = poll_health) -> dict:
    archive = Path(archive)
    _validate_identity(tag, commit, expected_sha256)
    _ensure_deployment_root(config, create=False)
    with _verified_archive_copy(archive, expected_sha256) as verified_archive:
        members, common_root = inspect_archive(verified_archive, config)
        final = config.releases / commit
        if final.exists() or final.is_symlink():
            raise DeploymentError(f"immutable release already exists: {commit}")
        previous = _current_commit(config)
        _ensure_releases_directory(config, create=True)
        pending = config.releases / f".pending-{commit}-{secrets.token_hex(6)}"
        if pending.exists() or pending.is_symlink():
            raise DeploymentError("pending release boundary already exists")
        pending.mkdir(mode=0o755)
        if pending.is_symlink() or not pending.is_dir():
            raise DeploymentError("pending release boundary is not a safe directory")
        try:
            _extract_validated(verified_archive, pending, members, common_root)
            metadata = {"schemaVersion": 1, "tag": tag, "commit": commit, "sha256": expected_sha256}
            (pending / ".timmy-release.json").write_text(json.dumps(metadata, sort_keys=True) + "\n", encoding="utf-8")
            os.replace(pending, final)
        except Exception:
            shutil.rmtree(pending, ignore_errors=True)
            raise
    _atomic_point(config, commit)
    try:
        _verify(config, commit, run_command, health_check, smoke=True)
    except Exception as error:
        _atomic_point(config, previous)
        try:
            if previous is not None:
                _verify(config, previous, run_command, health_check, smoke=False)
            else:
                run_command(config.stop_command, check=True, text=True, capture_output=True, shell=False, timeout=config.command_timeout)
        except Exception as rollback_error:
            raise DeploymentError("promotion verification failed and rollback verification also failed") from rollback_error
        if previous is None:
            raise DeploymentError("promotion verification failed; no prior release; service stopped") from error
        raise DeploymentError("promotion verification failed; prior release restored") from error
    return {"ok": True, "action": "promote", **metadata, "previousCommit": previous}


def _validate_release_directory(config: DeploymentConfig, commit: str) -> Path:
    release = config.releases / commit
    if not release.is_dir() or release.is_symlink():
        raise DeploymentError(f"release does not exist: {commit}")
    metadata_path = release / ".timmy-release.json"
    entrypoint = release / "server.mjs"
    if metadata_path.is_symlink() or entrypoint.is_symlink() or not entrypoint.is_file():
        raise DeploymentError("release metadata or entrypoint is invalid")
    try:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise DeploymentError("release metadata is missing or invalid") from error
    if metadata.get("commit") != commit or not TAG_RE.fullmatch(str(metadata.get("tag", ""))):
        raise DeploymentError("release metadata does not match requested commit")
    return release


def rollback(*, config: DeploymentConfig, commit: str, run_command: RunCommand = _run_argv,
             health_check: HealthCheck = poll_health) -> dict:
    if not COMMIT_RE.fullmatch(commit):
        raise DeploymentError("commit must be exactly 40 lowercase hexadecimal characters")
    _validate_release_directory(config, commit)
    previous = _current_commit(config)
    if previous == commit:
        raise DeploymentError("requested release is already current")
    _atomic_point(config, commit)
    try:
        _verify(config, commit, run_command, health_check, smoke=False)
    except Exception as error:
        _atomic_point(config, previous)
        if previous is not None:
            try:
                _verify(config, previous, run_command, health_check, smoke=False)
            except Exception as rollback_error:
                raise DeploymentError("rollback failed and prior release could not be restored") from rollback_error
        raise DeploymentError("rollback verification failed; prior release restored") from error
    return {"ok": True, "action": "rollback", "commit": commit, "previousCommit": previous}


def status(config: DeploymentConfig) -> dict:
    commit = _current_commit(config)
    if commit is None:
        return {"ok": True, "active": False, "commit": None, "release": None}
    metadata_path = config.releases / commit / ".timmy-release.json"
    metadata = {}
    try:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        pass
    return {"ok": True, "active": True, "commit": commit, "release": metadata.get("tag"), "sha256": metadata.get("sha256")}


def _json_argv(value: str, option: str) -> tuple[str, ...]:
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as error:
        raise DeploymentError(f"{option} must be a JSON argv array") from error
    if not isinstance(parsed, list) or not parsed or not all(isinstance(item, str) and item for item in parsed):
        raise DeploymentError(f"{option} must be a non-empty JSON argv array of strings")
    return tuple(parsed)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path(os.environ.get("TIMMY_STAGING_ROOT", "/opt/timmy-staging")))
    parser.add_argument("--dry-run", action="store_true", help="validate and print the intended action without mutation or commands")
    parser.add_argument("--restart-command", default=os.environ.get("TIMMY_STAGING_RESTART_COMMAND", '["systemctl","restart","timmy-staging.service"]'))
    parser.add_argument("--stop-command", default=os.environ.get("TIMMY_STAGING_STOP_COMMAND", '["systemctl","stop","timmy-staging.service"]'))
    parser.add_argument("--smoke-command", default=os.environ.get("TIMMY_STAGING_SMOKE_COMMAND", '["npm","run","test:staging-smoke"]'))
    parser.add_argument("--health-url", default=os.environ.get("TIMMY_STAGING_HEALTH_URL", "http://127.0.0.1:4174/timmy-staging/api/healthz"))
    subparsers = parser.add_subparsers(dest="command", required=True)
    promote_parser = subparsers.add_parser("promote")
    promote_parser.add_argument("--tag", required=True)
    promote_parser.add_argument("--archive", required=True, type=Path)
    promote_parser.add_argument("--sha256", required=True)
    promote_parser.add_argument("--commit", required=True)
    rollback_parser = subparsers.add_parser("rollback")
    rollback_parser.add_argument("--commit", required=True)
    subparsers.add_parser("status")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        config = DeploymentConfig(
            root=args.root,
            restart_command=_json_argv(args.restart_command, "restart command"),
            stop_command=_json_argv(args.stop_command, "stop command"),
            smoke_command=_json_argv(args.smoke_command, "smoke command"),
            health_url=args.health_url,
        )
        if args.command == "status":
            result = status(config)
        elif args.command == "promote":
            if args.dry_run:
                _validate_identity(args.tag, args.commit, args.sha256)
                _ensure_deployment_root(config, create=False)
                with _verified_archive_copy(args.archive, args.sha256) as verified_archive:
                    members, wrapper = inspect_archive(verified_archive, config)
                result = {"ok": True, "dryRun": True, "action": "promote", "commit": args.commit, "tag": args.tag, "members": len(members), "wrapper": wrapper}
            else:
                result = promote(config=config, tag=args.tag, archive=args.archive, expected_sha256=args.sha256, commit=args.commit)
        else:
            if args.dry_run:
                if not COMMIT_RE.fullmatch(args.commit):
                    raise DeploymentError("rollback release does not exist or commit is invalid")
                _validate_release_directory(config, args.commit)
                result = {"ok": True, "dryRun": True, "action": "rollback", "commit": args.commit}
            else:
                result = rollback(config=config, commit=args.commit)
        print(json.dumps(result, sort_keys=True))
        return 0
    except DeploymentError as error:
        print(f"deploy_staging: {error}", file=sys.stderr)
        return 2


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