#!/usr/bin/env python3
"""MAS evidence capture and restore rehearsal. No SMTP, no replica changes."""

from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import stat
import subprocess
import tarfile
import tempfile
from contextlib import closing
from pathlib import Path

FILES = ("events.jsonl", "tickets.sqlite3", "tickets-outbox.sqlite3")
PREFIX = "platform-pg/application-archives/mas/"
MAX_BYTES = 64 * 1024**2


class Refused(RuntimeError):
    pass


def digest(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def checkpoint(path, value):
    fd, name = tempfile.mkstemp(prefix="receipt-", dir=path.parent)
    with os.fdopen(fd, "w") as stream:
        json.dump(value, stream, indent=2)
        stream.flush()
        os.fsync(stream.fileno())
    os.chmod(name, 0o600)
    Path(name).replace(path)


def _regular(path):
    info = path.lstat()
    if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
        raise Refused("evidence_file_required")
    if info.st_size > MAX_BYTES or "mas-relay" in path.parts:
        raise Refused("evidence_path_outside_contract")


def _sqlite_fingerprint(path):
    with closing(sqlite3.connect(path.absolute().as_uri() + "?mode=ro", uri=True)) as database:
        database.execute("PRAGMA trusted_schema=OFF")
        if database.execute("PRAGMA integrity_check").fetchall() != [("ok",)]:
            raise Refused("sqlite_integrity_failed")
        tables = sorted(
            row[0]
            for row in database.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
            )
        )
        result = {"tables": {}}
        for name in tables:
            if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
                raise Refused("unexpected_sqlite_table")
            rows = [
                json.dumps(row, separators=(",", ":"), default=lambda value: value.hex()).encode()
                for row in database.execute(f'SELECT * FROM "{name}"')
            ]
            result["tables"][name] = {
                "rows": len(rows),
                "sha256": hashlib.sha256(b"\n".join(sorted(rows))).hexdigest(),
            }
        return result


def _backup_sqlite(source, output):
    _regular(source)
    if output.exists() or output.is_symlink():
        raise Refused("capture_destination_exists")
    fd = os.open(output, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    os.close(fd)
    with closing(sqlite3.connect(source.absolute().as_uri() + "?mode=ro", uri=True)) as src, closing(
        sqlite3.connect(output)
    ) as dst:
        src.backup(dst, pages=64)
    return {"sha256": digest(output), "bytes": output.stat().st_size, **_sqlite_fingerprint(output)}


def capture(source, destination, *, captured_at, image, writers_quiesced, replicas):
    """Copy the three evidence files. Caller must already have quiesced writers."""
    if writers_quiesced is not True or replicas != {"mas-mx": 1, "mas-statute": 1}:
        raise Refused("writers_not_quiesced")
    destination.mkdir(mode=0o700, exist_ok=False)
    files = {}
    for name in FILES:
        path = source / name
        if name == "tickets-outbox.sqlite3" and not path.exists():
            files[name] = {"present": False}
            continue
        if name.endswith(".sqlite3"):
            files[name] = {"present": True, **_backup_sqlite(path, destination / name)}
        else:
            _regular(path)
            target = destination / name
            fd = os.open(target, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
            with os.fdopen(fd, "wb") as out, path.open("rb") as src:
                while block := src.read(1024**2):
                    out.write(block)
            files[name] = {"present": True, "sha256": digest(target), "bytes": target.stat().st_size}
    receipt = {
        "status": "captured",
        "application": "mas",
        "captured_at": captured_at,
        "image": image,
        "writers_quiesced": True,
        "replicas": replicas,
        "mail_sent": False,
        "files": files,
    }
    checkpoint(destination / "receipt.json", receipt)
    return receipt


def seal(directory, stage, recipient):
    stage.mkdir(mode=0o700, exist_ok=True)
    encrypted = stage / "packet.tar.age"
    if encrypted.exists() or encrypted.is_symlink():
        raise Refused("partial_encryption_requires_review")
    archive = stage / "packet.tar"
    with archive.open("xb") as stream:
        os.chmod(archive, 0o600)
        with tarfile.open(fileobj=stream, mode="w") as packet:
            for name in (*FILES, "receipt.json"):
                path = directory / name
                if path.exists():
                    packet.add(path, arcname=name, recursive=False)
    subprocess.run(["age", "-r", recipient, "-o", str(encrypted), str(archive)], check=True, timeout=120)
    encrypted.chmod(0o600)
    archive.unlink()
    return {"ciphertext_sha256": digest(encrypted), "ciphertext_bytes": encrypted.stat().st_size}


def rehearse(encrypted, identity, destination):
    """Decrypt and compare hashes. Do not restore onto the facility or send mail."""
    if destination.exists() or destination.is_symlink():
        raise Refused("fresh_destination_required")
    destination.mkdir(mode=0o700)
    packet = destination / "packet.tar"
    result = subprocess.run(
        ["age", "-d", "-i", str(identity), "-o", str(packet), str(encrypted)],
        capture_output=True,
        timeout=120,
    )
    if result.returncode:
        raise Refused("decryption_failed")
    os.chmod(packet, 0o600)
    with tarfile.open(packet, "r:") as archive:
        entries = archive.getmembers()
        names = {entry.name for entry in entries}
        if (
            "receipt.json" not in names
            or any(not entry.isfile() or "/" in entry.name or entry.name.startswith(".") for entry in entries)
        ):
            raise Refused("invalid_packet_members")
        for entry in entries:
            target = destination / entry.name
            with archive.extractfile(entry) as source, target.open("xb") as out:
                os.chmod(target, 0o600)
                out.write(source.read())
    receipt = json.loads((destination / "receipt.json").read_text())
    if receipt.get("application") != "mas" or receipt.get("mail_sent") is not False:
        raise Refused("capture_receipt_rejected")
    for name, meta in receipt["files"].items():
        path = destination / name
        if not meta.get("present"):
            if path.exists():
                raise Refused("unexpected_absent_file")
            continue
        if digest(path) != meta["sha256"] or path.stat().st_size != meta["bytes"]:
            raise Refused("rehearsal_mismatch")
        path.unlink()
    packet.unlink()
    (destination / "receipt.json").unlink()
    proof = {
        "status": "rehearsal_matched",
        "application": "mas",
        "ciphertext_sha256": digest(encrypted),
        "mail_sent": False,
        "production_restored": False,
        "facility_scaled": False,
    }
    checkpoint(destination / "rehearsal.json", proof)
    return proof