#!/usr/bin/env python3
"""Node-only ledger capture, preview by default; execution briefly pauses email."""

from __future__ import annotations

import argparse
import copy
import fcntl
import json
import os
import secrets
import signal
import subprocess
import time
from pathlib import Path

from ledger_archive import MAX_BYTES, Refused, digest, fingerprint, private_directory, write_private

NS = "email-connect"
CLUSTER_UID = "a553c742-0115-43d4-99a4-a5ca56fe0786"
IMAGE = (
    "forgejo.coulomb.social/coulomb/email-connect@sha256:"
    "1e613d43734c44c5cd32c6083fda8b4da5d01ef43d10c6ca9c29a33c788cf75c"
)
LOCK_PATH = Path("/var/lib/email-connect-backup/capture.lock")
SELECTOR = "app.kubernetes.io/name=email-connect"


def k(*args, body=None, timeout=180):
    result = subprocess.run(["kubectl", *args], input=body, capture_output=True, timeout=timeout)
    if result.returncode:
        raise Refused("cluster_command_failed")
    return result.stdout


def document(*args, body=None):
    return json.loads(k(*args, "-o", "json", body=body, timeout=30))


def helper(name):
    return {
        "apiVersion": "v1",
        "kind": "Pod",
        "metadata": {"name": name, "namespace": NS, "labels": {"email-ledger-capture": name}},
        "spec": {
            "restartPolicy": "Never",
            "automountServiceAccountToken": False,
            "securityContext": {
                "runAsNonRoot": True,
                "runAsUser": 10001,
                "runAsGroup": 10001,
                "fsGroup": 10001,
                "fsGroupChangePolicy": "OnRootMismatch",
                "seccompProfile": {"type": "RuntimeDefault"},
            },
            "containers": [
                {
                    "name": "capture",
                    "image": IMAGE,
                    "command": ["python", "-c", "import time;time.sleep(1800)"],
                    "resources": {
                        "requests": {"cpu": "10m", "memory": "32Mi"},
                        "limits": {"cpu": "100m", "memory": "256Mi"},
                    },
                    "securityContext": {
                        "allowPrivilegeEscalation": False,
                        "readOnlyRootFilesystem": True,
                        "capabilities": {"drop": ["ALL"]},
                    },
                    "volumeMounts": [
                        {"name": "data", "mountPath": "/data", "readOnly": True},
                        {"name": "tmp", "mountPath": "/tmp"},
                    ],
                }
            ],
            "volumes": [
                {
                    "name": "data",
                    "persistentVolumeClaim": {"claimName": "email-connect-data", "readOnly": True},
                },
                {"name": "tmp", "emptyDir": {"sizeLimit": "160Mi"}},
            ],
        },
    }


def replica_patch(template, expected, desired):
    return json.dumps(
        [
            {"op": "test", "path": "/spec/template", "value": template},
            {"op": "test", "path": "/spec/replicas", "value": expected},
            {"op": "replace", "path": "/spec/replicas", "value": desired},
        ]
    ).encode()


def selected(selector, labels):
    if any(labels.get(key) != value for key, value in selector.get("matchLabels", {}).items()):
        return False
    for expression in selector.get("matchExpressions", []):
        key, op, values = expression["key"], expression["operator"], expression.get("values", [])
        if op == "In" and labels.get(key) not in values:
            return False
        if op == "NotIn" and labels.get(key) in values:
            return False
        if op == "Exists" and key not in labels:
            return False
        if op == "DoesNotExist" and key in labels:
            return False
        if op not in {"In", "NotIn", "Exists", "DoesNotExist"}:
            raise Refused("unknown_network_selector")
    return True


def require_no_helper_egress(policies, labels):
    applicable = [
        p["spec"]
        for p in policies
        if "Egress" in p["spec"].get("policyTypes", [])
        and selected(p["spec"].get("podSelector", {}), labels)
    ]
    if not applicable or any(spec.get("egress") for spec in applicable):
        raise Refused("helper_egress_not_denied")


def preflight():
    if document("get", "namespace", "kube-system")["metadata"]["uid"] != CLUSTER_UID:
        raise Refused("wrong_cluster")
    before = document("get", "deployment", "email-connect", "-n", NS)
    spec = before["spec"]
    containers = spec["template"]["spec"]["containers"]
    if (
        spec["replicas"] != 1
        or before["status"].get("readyReplicas") != 1
        or len(containers) != 1
        or containers[0]["image"] != IMAGE
    ):
        raise Refused("unexpected_release")
    if document("get", "pvc", "email-connect-data", "-n", NS)["status"]["phase"] != "Bound":
        raise Refused("pvc_not_bound")
    config = document("get", "configmap", "email-connect-config", "-n", NS)["data"]
    if config.get("EMAIL_CONNECT_DATABASE_PATH") != "/data/email-connect.db":
        raise Refused("unexpected_database_path")
    name = "email-ledger-capture-" + secrets.token_hex(8)
    pod = helper(name)
    policies = document("get", "networkpolicies", "-n", NS)["items"]
    require_no_helper_egress(policies, pod["metadata"]["labels"])
    body = json.dumps(pod).encode()
    k("create", "--dry-run=server", "-f", "-", body=body)
    projected = document(
        "patch",
        "deployment",
        "email-connect",
        "-n",
        NS,
        "--type=json",
        "--patch-file=/dev/stdin",
        "--dry-run=server",
        body=replica_patch(spec["template"], 1, 0),
    )
    expected = copy.deepcopy(spec)
    expected["replicas"] = 0
    if projected["spec"] != expected:
        raise Refused("unexpected_pause_spec")
    return before, body


def capture_database(name, destination):
    # No app startup, mounted secrets, SMTP import or provider request in helper.
    code = """import os,sqlite3,sys
source=sqlite3.connect('file:/data/email-connect.db?mode=ro',uri=True)
pages=source.execute('PRAGMA page_count').fetchone()[0]
size=source.execute('PRAGMA page_size').fetchone()[0]
assert pages*size<=LIMIT
os.umask(0o077)
output=sqlite3.connect('/tmp/ledger.sqlite3')
source.backup(output,pages=256)
assert output.execute('PRAGMA integrity_check').fetchall()==[('ok',)]
output.close();source.close()
with open('/tmp/ledger.sqlite3','rb') as stream:sys.stdout.buffer.write(stream.read())
""".replace("LIMIT", str(MAX_BYTES))
    data = k("exec", "-n", NS, name, "--", "python", "-c", code)
    if not data.startswith(b"SQLite format 3\x00") or len(data) > MAX_BYTES:
        raise Refused("invalid_ledger_backup")
    path = destination / "database.sqlite3"
    write_private(path, data)
    return {
        "database_sha256": digest(path),
        "database_bytes": len(data),
        "ledger_fingerprint": fingerprint(path),
        "sqlite_integrity": "ok",
        "writers_quiesced": True,
        "source_pvc_read_only": True,
    }


def capture(before, body, destination):
    if not destination.is_absolute() or destination.exists() or destination.is_symlink():
        raise Refused("fresh_absolute_destination_required")
    destination.mkdir(mode=0o700)
    private_directory(destination)
    receipt = {
        "status": "incomplete",
        "application": "email-connect",
        "image": IMAGE,
        "captured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    started, paused, created, captured = time.monotonic(), False, False, False
    name = json.loads(body)["metadata"]["name"]
    cleanup_error = None
    try:
        paused = True  # A failed response can still have changed replica count.
        k(
            "patch",
            "deployment",
            "email-connect",
            "-n",
            NS,
            "--type=json",
            "--patch-file=/dev/stdin",
            body=replica_patch(before["spec"]["template"], 1, 0),
        )
        k("wait", "--for=delete", "pod", "-n", NS, "-l", SELECTOR, "--timeout=120s")
        current = document("get", "deployment", "email-connect", "-n", NS)
        expected = copy.deepcopy(before["spec"])
        expected["replicas"] = 0
        if current["spec"] != expected:
            raise Refused("writers_not_quiesced")
        # Require no other pod mounting the ledger PVC, including a prior helper.
        pods = document("get", "pods", "-n", NS)["items"]
        if any(
            any(
                v.get("persistentVolumeClaim", {}).get("claimName") == "email-connect-data"
                for v in pod["spec"].get("volumes", [])
            )
            for pod in pods
        ):
            raise Refused("other_ledger_pvc_consumer")
        created = True
        k("create", "-f", "-", body=body)
        k("wait", "--for=condition=Ready", "pod/" + name, "-n", NS, "--timeout=120s")
        receipt.update(capture_database(name, destination))
        captured = True
    except BaseException as error:
        receipt.update(failure_type=type(error).__name__)
        raise
    finally:
        if created:
            try:
                k(
                    "delete",
                    "pod",
                    name,
                    "-n",
                    NS,
                    "--ignore-not-found=true",
                    "--wait=true",
                    "--timeout=120s",
                )
                receipt["helper_removed"] = True
            except BaseException as error:
                receipt["helper_cleanup_failed"] = True
                cleanup_error = error
        if paused:
            try:
                current = document("get", "deployment", "email-connect", "-n", NS)
                expected = copy.deepcopy(before["spec"])
                expected["replicas"] = current["spec"]["replicas"]
                if current["spec"] != expected or current["spec"]["replicas"] not in {0, 1}:
                    raise Refused("foreign_release_requires_receiver_check")
                if current["spec"]["replicas"] == 0:
                    k(
                        "patch",
                        "deployment",
                        "email-connect",
                        "-n",
                        NS,
                        "--type=json",
                        "--patch-file=/dev/stdin",
                        body=replica_patch(before["spec"]["template"], 0, 1),
                    )
                k(
                    "rollout",
                    "status",
                    "deployment/email-connect",
                    "-n",
                    NS,
                    "--timeout=180s",
                    timeout=240,
                )
                current = document("get", "deployment", "email-connect", "-n", NS)
                if current["spec"] != before["spec"] or current["status"].get("readyReplicas") != 1:
                    raise Refused("receiver_resume_not_verified")
                receipt["receiver_resumed"] = True
                receipt["unchanged_receiver_spec"] = True
            except BaseException as error:
                receipt["receiver_state_requires_check"] = True
                cleanup_error = cleanup_error or error
        receipt["status"] = (
            "captured"
            if captured and receipt.get("receiver_resumed") and not cleanup_error
            else "failed"
        )
        receipt["elapsed_seconds"] = round(time.monotonic() - started, 2)
        write_private(destination / "receipt.json", json.dumps(receipt, indent=2).encode())
        if cleanup_error:
            raise Refused("capture_cleanup_requires_check") from None
    return receipt


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--execute", action="store_true")
    parser.add_argument("--destination", type=Path, required=True)
    args = parser.parse_args()
    os.environ["KUBECONFIG"] = "/etc/rancher/k3s/k3s.yaml"
    before, body = preflight()
    if not args.execute:
        print(
            json.dumps(
                {
                    "status": "reviewed",
                    "live_changes": False,
                    "server_dry_run": True,
                    "image_pinned": True,
                    "pvc_read_only": True,
                    "receiver_spec_changed": False,
                }
            )
        )
        return 0
    LOCK_PATH.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    private_directory(LOCK_PATH.parent)
    fd = os.open(LOCK_PATH, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
    with os.fdopen(fd, "w") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise Refused("capture_already_running") from None
        # Recheck after locking; do not use a stale preflight across another run.
        before, body = preflight()
        print(json.dumps(capture(before, body, args.destination)))
    return 0


if __name__ == "__main__":

    def terminate(signum, frame):
        raise Refused("capture_interrupted")

    signal.signal(signal.SIGTERM, terminate)
    try:
        raise SystemExit(main())
    except Exception as error:
        print(json.dumps({"status": "failed", "failure_type": type(error).__name__}))
        raise SystemExit(1)