#!/usr/bin/env python3
"""Create and inspect atomic, age-encrypted Railiance S1 backup bundles."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import socket
import stat
import subprocess
import sys
import tarfile
import tempfile
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path, PurePosixPath
from typing import Any

from s1_backup_contract import (
    FORBIDDEN_PARTS,
    BackupContractError,
    declaration_sha256,
    load_spec,
)


ROOT = Path(__file__).resolve().parents[1]
BUNDLE_RE = re.compile(r"^s1-backup-(\d{8}T\d{6}Z)$")
REVISION_RE = re.compile(r"^[0-9a-f]{7,40}$")


class BackupError(RuntimeError):
    """A backup could not produce one complete, unambiguous bundle."""


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


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _rfc3339(value: datetime) -> str:
    return value.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _revision(explicit: str | None) -> str:
    value = explicit
    if value is None:
        try:
            if subprocess.check_output(
                ["git", "status", "--porcelain"],
                cwd=ROOT,
                text=True,
                stderr=subprocess.DEVNULL,
            ).strip():
                raise BackupError("source checkout is dirty; commit the exact backup implementation")
            value = subprocess.check_output(
                ["git", "rev-parse", "HEAD"], cwd=ROOT, text=True, stderr=subprocess.DEVNULL
            ).strip()
        except (OSError, subprocess.CalledProcessError) as exc:
            raise BackupError("source revision is unavailable; pass --source-revision") from exc
    if not REVISION_RE.fullmatch(value):
        raise BackupError("source revision must be a 7-40 character lowercase git revision")
    return value


def _source_path(source_root: Path, declared: str) -> Path:
    relative = PurePosixPath(declared).relative_to("/")
    return source_root.joinpath(*relative.parts)


def _walk_members(source_root: Path, declared: str) -> list[Path]:
    start = _source_path(source_root, declared)
    if not start.exists() and not start.is_symlink():
        return []
    members = [start]
    if start.is_dir() and not start.is_symlink():
        for current, directories, files in os.walk(start, followlinks=False):
            directories.sort()
            files.sort()
            base = Path(current)
            members.extend(base / name for name in directories)
            members.extend(base / name for name in files)
    return members


def _member_metadata(path: Path, source_root: Path) -> dict[str, Any]:
    relative = path.relative_to(source_root)
    lowered = {part.lower() for part in relative.parts}
    filename = relative.name.lower()
    if lowered & FORBIDDEN_PARTS or filename.endswith((".key", ".pem")) or any(
        marker in filename for marker in ("credential", "private_key", "token")
    ):
        raise BackupError(f"runtime member is secret/private-key shaped: /{relative}")
    info = path.lstat()
    if stat.S_ISREG(info.st_mode):
        kind = "file"
    elif stat.S_ISDIR(info.st_mode):
        kind = "directory"
    elif stat.S_ISLNK(info.st_mode):
        kind = "symlink"
        target = Path(os.readlink(path))
        resolved = target if target.is_absolute() else path.parent / target
        try:
            resolved.resolve(strict=False).relative_to(source_root)
        except ValueError as exc:
            raise BackupError(f"runtime link escapes fixture/source root: /{relative}") from exc
    else:
        raise BackupError(f"unsupported filesystem object: /{path.relative_to(source_root)}")
    return {
        "kind": kind,
        "mode": f"{stat.S_IMODE(info.st_mode):04o}",
        "path": str(PurePosixPath(*relative.parts)),
        "size_bytes": info.st_size if kind == "file" else 0,
    }


def _build_os_archive(
    *, source_root: Path, declaration: dict[str, Any], destination: Path
) -> list[dict[str, Any]]:
    selected: list[Path] = []
    missing_required: list[str] = []
    for entry in declaration["artifacts"]["os-config"]["paths"]:
        members = _walk_members(source_root, entry["path"])
        if not members:
            if entry["required"]:
                missing_required.append(entry["path"])
            continue
        selected.extend(members)
    if missing_required:
        raise BackupError("required backup inputs are missing: " + ", ".join(missing_required))
    unique = sorted(set(selected), key=lambda path: str(path.relative_to(source_root)))
    metadata = [_member_metadata(path, source_root) for path in unique]
    by_path = {item["path"]: item for item in metadata}
    for entry in declaration["artifacts"]["os-config"]["paths"]:
        if not entry["required"]:
            continue
        relative = str(PurePosixPath(entry["path"]).relative_to("/"))
        if by_path[relative]["mode"] != entry["expected_mode"]:
            raise BackupError(
                f"critical mode mismatch for {entry['path']}: "
                f"expected {entry['expected_mode']}, observed {by_path[relative]['mode']}"
            )
    with tarfile.open(destination, "w:gz", compresslevel=9) as archive:
        for path in unique:
            archive.add(
                path,
                arcname=str(PurePosixPath(*path.relative_to(source_root).parts)),
                recursive=False,
            )
    return metadata


def _write_packages(
    destination: Path, declaration: dict[str, Any], packages_file: Path | None
) -> None:
    if packages_file is not None:
        if not packages_file.is_file():
            raise BackupError("--packages-file must identify a readable regular file")
        shutil.copyfile(packages_file, destination)
        return
    command = declaration["artifacts"]["packages"]["command"]
    try:
        with destination.open("wb") as output:
            subprocess.run(command, stdout=output, stderr=subprocess.DEVNULL, check=True)
    except (OSError, subprocess.CalledProcessError) as exc:
        raise BackupError("fixed package-selection command failed") from exc


def _age_binary(explicit: str | None) -> str:
    if explicit:
        return explicit
    for candidate in (shutil.which("age"), "/usr/local/bin/age"):
        if candidate and Path(candidate).is_file():
            return candidate
    raise BackupError("age executable not found")


def _encrypt(age_binary: str, recipients: list[str], source: Path, destination: Path) -> None:
    command = [age_binary]
    for recipient in recipients:
        command.extend(["-r", recipient])
    command.extend(["-o", str(destination), str(source)])
    try:
        subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
    except (OSError, subprocess.CalledProcessError) as exc:
        raise BackupError("age encryption failed; no bundle was published") from exc
    if not destination.is_file() or destination.stat().st_size == 0:
        raise BackupError("age encryption produced no usable output")


def _bundle_size(path: Path) -> int:
    return sum(item.stat().st_size for item in path.rglob("*") if item.is_file())


def _complete_bundles(output_dir: Path) -> list[Path]:
    bundles = []
    if not output_dir.is_dir():
        return bundles
    for path in output_dir.iterdir():
        if not path.is_dir() or not BUNDLE_RE.fullmatch(path.name):
            continue
        receipt = path / "receipt.json"
        try:
            payload = json.loads(receipt.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        if payload.get("event_type") == "backup" and payload.get("status") == "pass":
            bundles.append(path)
    return sorted(bundles, key=lambda item: item.name, reverse=True)


def _retention_candidates(output_dir: Path, keep: int, max_bytes: int) -> list[Path]:
    bundles = _complete_bundles(output_dir)
    retained = list(bundles)
    candidates: list[Path] = []
    while len(retained) > keep:
        candidates.append(retained.pop())
    while sum(_bundle_size(path) for path in retained) > max_bytes and len(retained) > 2:
        candidates.append(retained.pop())
    return candidates


def _prune_approval(output_dir: Path, candidates: list[Path]) -> str:
    material = "\n".join(
        f"{path.name}:{_sha256(path / 'receipt.json')}" for path in sorted(candidates)
    )
    digest = hashlib.sha256(material.encode()).hexdigest()[:16].upper()
    return f"PRUNE-S1-BACKUPS-{digest}"


def plan_prune(spec_path: Path, output_dir: Path | None) -> dict[str, Any]:
    declaration = load_spec(spec_path)
    root = Path(output_dir or declaration["backup_root"]).resolve()
    if root != Path(declaration["backup_root"]).resolve() and output_dir is None:
        raise BackupError("retention root does not match the declaration")
    retention = declaration["retention"]
    candidates = _retention_candidates(
        root,
        retention["keep_complete_bundles"],
        retention["max_total_bytes"],
    )
    return {
        "approval": _prune_approval(root, candidates) if candidates else None,
        "backup_root": str(root),
        "candidates": [path.name for path in candidates],
        "status": "approval-required" if candidates else "nothing-to-prune",
    }


def apply_prune(spec_path: Path, output_dir: Path | None, approval: str) -> dict[str, Any]:
    plan = plan_prune(spec_path, output_dir)
    if not plan["candidates"]:
        return plan
    if approval != plan["approval"]:
        raise BackupError("prune approval does not match the current exact candidate set")
    root = Path(plan["backup_root"])
    quarantine = root / f".prune-{uuid.uuid4()}"
    quarantine.mkdir(mode=0o700)
    moved: list[Path] = []
    try:
        for name in plan["candidates"]:
            source = root / name
            if source.parent != root or not BUNDLE_RE.fullmatch(source.name):
                raise BackupError("prune candidate escaped the declared backup root")
            target = quarantine / source.name
            os.replace(source, target)
            moved.append(target)
        directory_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
    except Exception:
        for target in reversed(moved):
            if target.exists():
                os.replace(target, root / target.name)
        shutil.rmtree(quarantine, ignore_errors=True)
        raise
    shutil.rmtree(quarantine)
    return {
        "backup_root": str(root),
        "deleted": plan["candidates"],
        "status": "pruned",
    }


def _retention_decision(output_dir: Path, keep: int, max_bytes: int, new_size: int) -> dict[str, Any]:
    projected = [(path.name, _bundle_size(path)) for path in _complete_bundles(output_dir)]
    projected.insert(0, ("new-bundle", new_size))
    candidates: list[str] = []
    total = sum(size for _, size in projected)
    retained = list(projected)
    while len(retained) > keep:
        name, size = retained.pop()
        candidates.append(name)
        total -= size
    while total > max_bytes and len(retained) > 2:
        name, size = retained.pop()
        candidates.append(name)
        total -= size
    return {
        "keep_complete_bundles": keep,
        "max_total_bytes": max_bytes,
        "projected_total_bytes": total,
        "prune_candidates": sorted(name for name in candidates if name != "new-bundle"),
        "pruning_applied": False,
    }


def create_bundle(
    *,
    spec_path: Path,
    source_root: Path,
    output_dir: Path | None,
    packages_file: Path | None,
    age_binary: str | None,
    source_revision: str | None,
    timestamp: str | None,
    hostname: str | None,
) -> Path:
    declaration = load_spec(spec_path)
    source_root = source_root.resolve()
    production = source_root == Path("/")
    if production and os.geteuid() != 0:
        raise BackupError("production backup requires root before creating the destination")
    if packages_file is not None and production:
        raise BackupError("--packages-file is allowed only with a fixture --source-root")
    destination_root = Path(output_dir or declaration["backup_root"])
    if not destination_root.is_absolute():
        raise BackupError("output directory must be absolute")
    if destination_root.resolve() in {Path("/"), Path("/etc"), Path("/home"), Path("/root")}:
        raise BackupError("output directory must be bounded")

    observed = _utc_now()
    stamp = timestamp or observed.strftime("%Y%m%dT%H%M%SZ")
    if not re.fullmatch(r"\d{8}T\d{6}Z", stamp):
        raise BackupError("timestamp must be YYYYMMDDTHHMMSSZ")
    bundle_name = f"s1-backup-{stamp}"
    final = destination_root / bundle_name
    if final.exists():
        raise BackupError(f"refusing to overwrite existing bundle {bundle_name}")

    destination_root.mkdir(parents=True, exist_ok=True, mode=0o700)
    stage = Path(tempfile.mkdtemp(prefix=f".{bundle_name}.", dir=destination_root))
    os.chmod(stage, 0o700)
    try:
        plain_os = stage / ".os-config.tar.gz"
        plain_packages = stage / ".packages.txt"
        members = _build_os_archive(
            source_root=source_root, declaration=declaration, destination=plain_os
        )
        _write_packages(plain_packages, declaration, packages_file)
        executable = _age_binary(age_binary)
        encrypted = {
            "os-config": stage / "os-config.tar.gz.age",
            "packages": stage / "packages.txt.age",
        }
        _encrypt(executable, declaration["recipients"], plain_os, encrypted["os-config"])
        _encrypt(executable, declaration["recipients"], plain_packages, encrypted["packages"])
        plain_os.unlink()
        plain_packages.unlink()

        revision = _revision(source_revision)
        declaration_digest = declaration_sha256(spec_path)
        artifacts = [
            {
                "artifact_class": artifact_class,
                "encrypted_sha256": _sha256(path),
                "filename": path.name,
                "size_bytes": path.stat().st_size,
            }
            for artifact_class, path in encrypted.items()
        ]
        manifest = {
            "schema_version": "1.0",
            "bundle_id": bundle_name,
            "created_at": _rfc3339(observed),
            "host": hostname or socket.gethostname(),
            "source_revision": revision,
            "declaration_sha256": declaration_digest,
            "recipient_sha256": [
                hashlib.sha256(value.encode()).hexdigest() for value in declaration["recipients"]
            ],
            "artifacts": artifacts,
            "os_config_members": members,
            "package_source": "fixture-file" if packages_file else "dpkg --get-selections",
        }
        manifest_path = stage / "manifest.json"
        manifest_path.write_text(
            json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        retention = declaration["retention"]
        receipt = {
            "schema_version": "1.0",
            "receipt_id": str(uuid.uuid4()),
            "event_type": "backup",
            "status": "pass",
            "created_at": _rfc3339(observed),
            "bundle_id": bundle_name,
            "host": manifest["host"],
            "source_revision": revision,
            "declaration_sha256": declaration_digest,
            "manifest_sha256": _sha256(manifest_path),
            "artifacts": artifacts,
            "retention": _retention_decision(
                destination_root,
                retention["keep_complete_bundles"],
                retention["max_total_bytes"],
                _bundle_size(stage),
            ),
        }
        (stage / "receipt.json").write_text(
            json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        for path in stage.iterdir():
            if path.is_file():
                with path.open("rb") as handle:
                    os.fsync(handle.fileno())
        os.replace(stage, final)
        directory_fd = os.open(destination_root, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
        return final
    except Exception:
        shutil.rmtree(stage, ignore_errors=True)
        raise


def inspect_status(spec_path: Path, output_dir: Path | None) -> dict[str, Any]:
    declaration = load_spec(spec_path)
    root = Path(output_dir or declaration["backup_root"])
    bundles = _complete_bundles(root)
    if not bundles:
        raise BackupError("no complete S1 backup bundle exists")
    newest = bundles[0]
    # Import lazily to keep the create path independent while making status
    # prove the newest evidence, rather than trusting a pass-shaped JSON file.
    from s1_restore import RestoreError, validate_bundle

    try:
        receipt = validate_bundle(newest, spec_path)["receipt"]
        created = datetime.fromisoformat(receipt["created_at"].replace("Z", "+00:00"))
    except (RestoreError, KeyError, TypeError, ValueError) as exc:
        raise BackupError(f"newest backup evidence is invalid: {exc}") from exc
    fresh_until = created + timedelta(hours=declaration["retention"]["freshness_hours"])
    total = sum(_bundle_size(bundle) for bundle in bundles)
    result = {
        "bundle": newest.name,
        "bundle_count": len(bundles),
        "created_at": receipt["created_at"],
        "fresh_until": _rfc3339(fresh_until),
        "status": "pass",
        "total_bytes": total,
    }
    if _utc_now() >= fresh_until:
        result["status"] = "stale"
    if len(bundles) > declaration["retention"]["keep_complete_bundles"]:
        result["status"] = "over-retention"
    if total > declaration["retention"]["max_total_bytes"]:
        result["status"] = "over-budget"
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    check = subparsers.add_parser("check", help="validate source-only backup inputs")
    check.add_argument("--spec", type=Path, default=ROOT / "spec" / "s1-backup.yaml")
    create = subparsers.add_parser("create", help="create one complete encrypted bundle")
    create.add_argument("--spec", type=Path, default=ROOT / "spec" / "s1-backup.yaml")
    create.add_argument("--source-root", type=Path, default=Path("/"))
    create.add_argument("--output-dir", type=Path)
    create.add_argument("--packages-file", type=Path)
    create.add_argument("--age-bin")
    create.add_argument("--source-revision")
    create.add_argument("--timestamp")
    create.add_argument("--hostname")
    status_parser = subparsers.add_parser("status", help="fail if backup evidence is absent or stale")
    status_parser.add_argument("--spec", type=Path, default=ROOT / "spec" / "s1-backup.yaml")
    status_parser.add_argument("--output-dir", type=Path)
    prune_plan = subparsers.add_parser(
        "prune-plan", help="print the exact retained bundles requiring deletion approval"
    )
    prune_plan.add_argument("--spec", type=Path, default=ROOT / "spec" / "s1-backup.yaml")
    prune_plan.add_argument("--output-dir", type=Path)
    prune = subparsers.add_parser("prune", help="apply one exact, freshly rendered prune plan")
    prune.add_argument("--spec", type=Path, default=ROOT / "spec" / "s1-backup.yaml")
    prune.add_argument("--output-dir", type=Path)
    prune.add_argument("--approval", required=True)
    args = parser.parse_args()
    try:
        if args.command == "check":
            declaration = load_spec(args.spec)
            result = {"ok": True, "artifacts": sorted(declaration["artifacts"])}
        elif args.command == "create":
            bundle = create_bundle(
                spec_path=args.spec,
                source_root=args.source_root,
                output_dir=args.output_dir,
                packages_file=args.packages_file,
                age_binary=args.age_bin,
                source_revision=args.source_revision,
                timestamp=args.timestamp,
                hostname=args.hostname,
            )
            result = {"bundle": str(bundle), "status": "pass"}
        elif args.command == "status":
            result = inspect_status(args.spec, args.output_dir)
        elif args.command == "prune-plan":
            result = plan_prune(args.spec, args.output_dir)
        else:
            result = apply_prune(args.spec, args.output_dir, args.approval)
        print(json.dumps(result, sort_keys=True))
        return 0 if result.get("status", "pass") == "pass" else 1
    except (BackupContractError, BackupError, OSError, tarfile.TarError) as exc:
        print(f"S1 backup failed closed: {exc}", file=sys.stderr)
        return 1


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