src/services/brain_backup/manager.py

"""Startup restore gate and conflict-safe content checkpoint coordinator."""
import hashlib
import json
import os
import re
import shutil
import tempfile
from datetime import datetime, timezone
from pathlib import Path

from .git import BackupError
from .postgres import canonical
from .state import generation, locked, write_status

CHUNK = 32 * 1024 * 1024


def sha(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def validate(root, profile, postgres_major=15):
    """Validate every payload before any database write. No extra source files."""
    from src.constants.data import Brain, Workspace
    from src.services.brain_identity import validate_brain_id
    root = Path(root)
    try:
        manifest = json.loads((root / "manifest.json").read_text())
        if (manifest["format_version"] != 1 or manifest["postgres_major"] != postgres_major
                or manifest["schema_revision"] != 1 or manifest["embedding"] != profile
                or manifest["stored_dimension"] != profile["stored_dimension"]
                or manifest["embedding_profile"] != profile["profile"]):
            raise BackupError("BACKUP_INCOMPATIBLE_MANIFEST")
        checks = manifest["checksums"]
        hold = manifest.get("maintenance", {})
        if hold and (set(hold) != {"id", "quiesced"} or hold["quiesced"] is not True
                     or not re.fullmatch(r"[0-9a-f]{32}", str(hold["id"]))):
            raise BackupError("BACKUP_INVALID_MAINTENANCE")
        actual = {p.relative_to(root).as_posix() for p in root.rglob("*")
                  if p.is_file() and p.relative_to(root).parts[0] != ".git"}
        expected = {"manifest.json", *checks}
        if actual != expected or any(p.is_symlink() for p in root.rglob("*") if p.relative_to(root).parts[0] != ".git"):
            raise BackupError("BACKUP_UNEXPECTED_FILES")
        for name, checksum in checks.items():
            if (not re.fullmatch(r"(?:system/registry\.json|brains/[A-Za-z0-9_-]+\.dump\.part[0-9]{4})", name)
                    or not re.fullmatch(r"[0-9a-f]{64}", checksum) or sha(root / name) != checksum):
                raise BackupError("BACKUP_CHECKSUM_MISMATCH")
        registry = json.loads((root / "system/registry.json").read_text())
        names, ids = set(), set()
        for b in registry["brains"]:
            validate_brain_id(b["name_key"])
            clean = Brain(id=b["id"], name_key=b["name_key"]).model_dump(mode="json")
            if b["document"] != clean or b["name_key"] in names or b["id"] in ids:
                raise BackupError("BACKUP_INVALID_REGISTRY")
            names.add(b["name_key"])
            ids.add(b["id"])
        if manifest["brains"] != sorted(names) or set(manifest["fingerprints"]) != names:
            raise BackupError("BACKUP_INVALID_REGISTRY")
        for w in registry["workspaces"]:
            obj = Workspace.model_validate(w["document"])
            if (obj.model_dump(mode="json") != w["document"] or obj.brain_id not in names
                    or (obj.id, obj.brain_id, obj.slug) != (w["id"], w["brain_id"], w["slug"])):
                raise BackupError("BACKUP_INVALID_WORKSPACE")
        payloads = {"system/registry.json"}
        for name, parts in manifest["archives"].items():
            if name not in names or not parts or parts != [f"brains/{name}.dump.part{i:04d}" for i in range(len(parts))]:
                raise BackupError("BACKUP_INVALID_ARCHIVE")
            payloads.update(parts)
        if payloads != set(checks):
            raise BackupError("BACKUP_INVALID_ARCHIVE")
        digest = hashlib.sha256(canonical({"registry": registry, "fingerprints": manifest["fingerprints"], "embedding": profile})).hexdigest()
        if digest != manifest["content_digest"]:
            raise BackupError("BACKUP_CONTENT_DIGEST_MISMATCH")
        return manifest, registry
    except BackupError:
        raise
    except (OSError, ValueError, KeyError, TypeError, AttributeError):
        raise BackupError("BACKUP_INVALID_MANIFEST") from None


class Manager:
    def __init__(self, repository, postgres, root):
        self.repo = repository
        self.pg = postgres
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        self.receipt_path = self.root / "receipt.json"

    def receipt(self):
        if not self.receipt_path.exists():
            return None
        try:
            receipt = json.loads(self.receipt_path.read_text())
            if not isinstance(receipt, dict) or set(receipt) != {"head", "digest"}:
                raise ValueError("invalid receipt")
            for key, length in (("head", 40), ("digest", 64)):
                if receipt[key] is not None and not re.fullmatch(f"[0-9a-f]{{{length}}}", receipt[key]):
                    raise ValueError("invalid receipt")
            return receipt
        except (OSError, ValueError):
            raise BackupError("BACKUP_LOCAL_RECEIPT_CORRUPT") from None

    def remember(self, head, digest):
        temp = self.receipt_path.with_suffix(".tmp")
        temp.write_bytes(canonical({"head": head, "digest": digest}))
        temp.replace(self.receipt_path)

    def remote(self):
        head = self.repo.fetch()
        if head is None:
            return None, None, None
        if not (self.repo.path / "manifest.json").exists():
            # An initialized README-only repository has no snapshot yet. A
            # deleted manifest or any backup-shaped content is corruption.
            files = set(filter(None, self.repo.run("ls-files", "-z").stdout.decode().split("\0")))
            history = self.repo.run("log", "-1", "--format=%H", "HEAD", "--", "manifest.json").stdout.strip()
            if history or files - {"README.md", "LICENSE", ".gitignore", ".gitattributes"}:
                raise BackupError("BACKUP_REMOTE_MANIFEST_MISSING")
            return head, None, None
        manifest, registry = validate(self.repo.path, self.pg.profile, self.pg.major_version())
        return head, manifest, registry

    def startup(self):
        with locked(self.root / "operation.lock"):
            head, manifest, registry = self.remote()
            local = self.pg.has_local_content()
            receipt = self.receipt()
            if not manifest and receipt and receipt["digest"]:
                raise BackupError("BACKUP_REMOTE_HISTORY_MISSING")
            if local:
                # Only a matching checkpoint lineage permits retaining local
                # content. Unrelated populated stores require explicit recovery.
                if manifest and (not receipt or receipt["head"] != head):
                    raise BackupError("BACKUP_LOCAL_REMOTE_CONFLICT")
            elif manifest:
                with tempfile.TemporaryDirectory(dir=self.root) as temp:
                    staged = Path(temp)
                    (staged / "brains").mkdir()
                    for name, parts in manifest["archives"].items():
                        archive = staged / "brains" / f"{name}.dump"
                        with archive.open("wb") as dest:
                            for part in parts:
                                with (self.repo.path / part).open("rb") as src:
                                    shutil.copyfileobj(src, dest)
                        self.pg.tool("pg_restore", self.pg.system, "--list", str(archive))
                    self.pg.restore(staged, registry)
            # A previous container may still be completing a shutdown push.
            # Never publish readiness/writability from a stale restore.
            latest, _, _ = self.remote()
            if latest != head:
                from .maintenance import fail
                fail("BACKUP_STARTUP_REMOTE_ADVANCED")
                raise BackupError("BACKUP_STARTUP_REMOTE_ADVANCED")
            from .maintenance import initialize
            initialize(manifest.get("maintenance") if manifest else None, preserve_work=local)
            self.remember(head, manifest["content_digest"] if manifest else None)
            write_status(state="ready", commit=head, restored=bool(manifest and not local),
                         restored_checkpoint=head if manifest and not local else None,
                         remote_checkpoint=head, restore_qualified=bool(manifest),
                         checkpoint_generation=generation(), remote_verified=True,
                         last_success=datetime.now(timezone.utc).isoformat() if manifest else None)

    def verify_checkpoint(self, expected):
        with locked(self.root / "operation.lock"):
            head, manifest, _ = self.remote()
            if head != expected or manifest is None:
                write_status(state="degraded", remote_checkpoint=head, remote_verified=False,
                             error="BACKUP_REMOTE_CONFLICT")
                raise BackupError("BACKUP_REMOTE_CONFLICT")
            write_status(remote_checkpoint=head, remote_verified=True)
            return head

    def verify_startup(self):
        receipt = self.receipt()
        if not receipt:
            raise BackupError("BACKUP_STARTUP_RESTORE_REQUIRED")
        with locked(self.root / "operation.lock"):
            head, _, _ = self.remote()
            if head != receipt["head"]:
                from .maintenance import fail
                fail("BACKUP_STARTUP_REMOTE_ADVANCED")
                raise BackupError("BACKUP_STARTUP_REMOTE_ADVANCED")
            write_status(remote_checkpoint=head, remote_verified=True)

    def reconcile(self):
        """A sealed handoff pins publication until explicit checkpoint/resume.

        Read paths can lazily initialize derived indexes. Neither periodic nor
        shutdown reconciliation may move the handoff behind the operator's SHA.
        """
        from .maintenance import status
        control = status()
        if control.get("quiesced") and control.get("sealed"):
            receipt = self.receipt()
            if not receipt or not receipt["head"]:
                raise BackupError("BACKUP_STARTUP_RESTORE_REQUIRED")
            self.verify_checkpoint(receipt["head"])
            return receipt["head"], False
        return self.backup()

    def backup(self, maintenance=None):
        with locked(self.root / "operation.lock"):
            head, manifest, _ = self.remote()
            receipt = self.receipt()
            if receipt is None or receipt["head"] != head:
                raise BackupError("BACKUP_REMOTE_CONFLICT")
            before = generation()
            if maintenance is None:
                from .maintenance import hold_manifest
                maintenance = hold_manifest()
            with tempfile.TemporaryDirectory(dir=self.root) as temp:
                output = Path(temp)
                registry, fingerprints = self.pg.snapshot(output)
                digest = hashlib.sha256(canonical({"registry": registry, "fingerprints": fingerprints,
                                                   "embedding": self.pg.profile})).hexdigest()
                if manifest and digest == manifest["content_digest"] and manifest.get("maintenance", {}) == maintenance:
                    write_status(state="ready", commit=head, dirty=generation() != before,
                                 checkpoint_generation=before, remote_checkpoint=head, remote_verified=True,
                                 last_success=datetime.now(timezone.utc).isoformat(),
                                 checked_at=datetime.now(timezone.utc).isoformat())
                    return head, False
                archives = {}
                for archive in (output / "brains").glob("*.dump"):
                    parts = []
                    with archive.open("rb") as src:
                        for i, chunk in enumerate(iter(lambda: src.read(CHUNK), b"")):
                            part = archive.with_name(f"{archive.name}.part{i:04d}")
                            part.write_bytes(chunk)
                            parts.append(part.relative_to(output).as_posix())
                    archives[archive.stem] = parts
                    archive.unlink()
                checks = {p.relative_to(output).as_posix(): sha(p) for p in output.rglob("*") if p.is_file()}
                from src.constants.data import BRAIN_VERSION
                document = {"format_version": 1, "schema_revision": 1, "postgres_major": self.pg.major_version(), "created_at": datetime.now(timezone.utc).isoformat(),
                            "brainapi_version": BRAIN_VERSION, "source_sha": (os.getenv("BUILD_SHA") or "unknown"),
                            "embedding_profile": self.pg.profile["profile"], "stored_dimension": self.pg.profile["stored_dimension"],
                            "embedding": self.pg.profile, "brains": sorted(fingerprints), "fingerprints": fingerprints,
                            "archives": archives, "checksums": checks, "content_digest": digest}
                document["maintenance"] = maintenance
                (output / "manifest.json").write_bytes(canonical(document))
                validate(output, self.pg.profile, self.pg.major_version())
                new_head = self.repo.publish(output, head)
                # A successful push alone is insufficient for a handoff receipt.
                verified, verified_manifest, _ = self.remote()
                if verified != new_head or verified_manifest is None:
                    raise BackupError("BACKUP_REMOTE_CONFLICT")
                self.remember(new_head, digest)
                write_status(state="ready", commit=new_head, dirty=generation() != before,
                             checkpoint_generation=before, remote_checkpoint=new_head,
                             remote_verified=True,
                             last_success=datetime.now(timezone.utc).isoformat())
                return new_head, True