#!/usr/bin/env python3
"""Capture one pg_dump and parity manifest from the same exported VPS snapshot."""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import re
import shlex
import shutil
import subprocess
import tarfile
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

try:
    from .private_receipt_io import print_private_receipt_summary, write_private_json
    from .verify_postgres_parity_manifest import load_manifest
except ImportError:  # pragma: no cover - direct script execution
    from private_receipt_io import print_private_receipt_summary, write_private_json
    from verify_postgres_parity_manifest import load_manifest


SAFE_NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}\Z")
SAFE_RUN_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{7,63}\Z")
SAFE_AUTHORIZATION_REF_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.:/@+-]{7,255}\Z")
SAFE_SSH_TARGET_RE = re.compile(r"(?:[A-Za-z0-9][A-Za-z0-9._-]{0,63}@)?[A-Za-z0-9][A-Za-z0-9.-]{0,252}\Z")
SNAPSHOT_ID_RE = re.compile(r"[0-9A-F]+-[0-9A-F]+-[0-9]+\Z")
TXID_SNAPSHOT_RE = re.compile(r"[0-9]+:[0-9]+:(?:[0-9]+(?:,[0-9]+)*)?\Z")
WAL_LSN_RE = re.compile(r"[0-9A-F]+/[0-9A-F]+\Z")
SYSTEM_IDENTIFIER_RE = re.compile(r"[0-9]+\Z")
SOURCE_SNAPSHOT_RECEIPT_SCHEMA = "livingip.sourceSnapshotReceipt.v2"
EXPECTED_FILES = {
    "teleo-canonical.dump",
    "teleo-canonical.dump.sha256",
    "teleo-canonical.toc",
    "source-manifest.jsonl",
    "source-context.txt",
    "source-snapshot.json",
    "service-before.txt",
    "service-after.txt",
}


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 sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def canonical_sha256(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
    return sha256_bytes(payload.encode())


def load_snapshot_metadata(path: Path) -> dict[str, str]:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise RuntimeError("source snapshot metadata is unreadable") from exc
    if not isinstance(payload, dict):
        raise RuntimeError("source snapshot metadata must be an object")
    fields = {
        "exported_snapshot_id": SNAPSHOT_ID_RE,
        "txid_snapshot": TXID_SNAPSHOT_RE,
        "wal_lsn": WAL_LSN_RE,
        "system_identifier": SYSTEM_IDENTIFIER_RE,
    }
    normalized: dict[str, str] = {}
    for field, pattern in fields.items():
        value = str(payload.get(field) or "")
        if not pattern.fullmatch(value):
            raise RuntimeError(f"source snapshot metadata has invalid {field}")
        normalized[field] = value
    captured_at_utc = str(payload.get("captured_at_utc") or "")
    if not captured_at_utc or any(character in captured_at_utc for character in "\r\n\x00"):
        raise RuntimeError("source snapshot metadata has invalid captured_at_utc")
    normalized["captured_at_utc"] = captured_at_utc
    return normalized


def load_service_state(path: Path) -> dict[str, str]:
    try:
        state = dict(line.split("=", 1) for line in path.read_text(encoding="utf-8").splitlines() if "=" in line)
        main_pid = int(state.get("MainPID") or 0)
        restarts = int(state.get("NRestarts") or 0)
    except (OSError, TypeError, ValueError) as exc:
        raise RuntimeError("source service state is unreadable") from exc
    if state.get("ActiveState") != "active" or state.get("SubState") != "running" or main_pid <= 0 or restarts < 0:
        raise RuntimeError("source service is not active/running with a positive MainPID")
    return state


def build_provenance_binding(
    *,
    run_id: str,
    authorization_ref: str,
    source: dict[str, str],
    snapshot: dict[str, str],
    dump_sha256: str,
    manifest_sql_sha256: str,
    source_manifest_sha256: str,
    source_context_sha256: str,
) -> dict[str, Any]:
    payload = {
        "receipt_schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
        "run_id": run_id,
        "capture_authorization_ref": authorization_ref,
        "source": source,
        "exported_snapshot_id": snapshot["exported_snapshot_id"],
        "txid_snapshot": snapshot["txid_snapshot"],
        "wal_lsn": snapshot["wal_lsn"],
        "source_system_identifier": snapshot["system_identifier"],
        "dump_sha256": dump_sha256,
        "manifest_sql_sha256": manifest_sql_sha256,
        "source_manifest_sha256": source_manifest_sha256,
        "source_context_sha256": source_context_sha256,
    }
    return {"algorithm": "sha256", "payload": payload, "sha256": canonical_sha256(payload)}


def ssh_base(identity_file: Path) -> list[str]:
    return [
        "ssh",
        "-T",
        "-i",
        str(identity_file),
        "-o",
        "BatchMode=yes",
        "-o",
        "ConnectTimeout=12",
    ]


def remote_capture_script(*, run_id: str, container: str, database: str, service: str) -> str:
    remote_root = f"/tmp/teleo-canonical-snapshot-{run_id}"
    container_manifest = f"/tmp/teleo-postgres-parity-{run_id}.sql"
    container_dump = f"/tmp/teleo-canonical-{run_id}.dump"
    quoted = {
        "remote_root": shlex.quote(remote_root),
        "container": shlex.quote(container),
        "database": shlex.quote(database),
        "service": shlex.quote(service),
        "container_manifest": shlex.quote(container_manifest),
        "container_dump": shlex.quote(container_dump),
    }
    return f"""set -euo pipefail
remote_root={quoted["remote_root"]}
container={quoted["container"]}
database={quoted["database"]}
service={quoted["service"]}
container_manifest={quoted["container_manifest"]}
container_dump={quoted["container_dump"]}

cleanup() {{
  docker exec "$container" rm -f "$container_manifest" "$container_dump" >/dev/null 2>&1 || true
  rm -rf "$remote_root"
}}
trap cleanup EXIT

assert_service_healthy() {{
  local state_file="$1"
  grep -qx 'ActiveState=active' "$state_file"
  grep -qx 'SubState=running' "$state_file"
  local main_pid restarts
  main_pid="$(sed -n 's/^MainPID=//p' "$state_file")"
  restarts="$(sed -n 's/^NRestarts=//p' "$state_file")"
  [[ "$main_pid" =~ ^[1-9][0-9]*$ ]]
  [[ "$restarts" =~ ^[0-9]+$ ]]
}}

test -d "$remote_root"
test -f "$remote_root/postgres-parity-manifest.sql"
docker inspect "$container" >/dev/null
docker exec "$container" pg_isready -U postgres -d "$database" >/dev/null
systemctl show "$service" -p ActiveState -p SubState -p MainPID -p NRestarts -p ActiveEnterTimestamp --no-pager > "$remote_root/service-before.txt"
assert_service_healthy "$remote_root/service-before.txt"
docker inspect "$container" --format 'container_id={{{{.Id}}}}\nrunning={{{{.State.Running}}}}\nnetwork_mode={{{{.HostConfig.NetworkMode}}}}' > "$remote_root/source-context.txt"
printf 'database=%s\ncaptured_at_utc=%s\n' "$database" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" >> "$remote_root/source-context.txt"
docker cp "$remote_root/postgres-parity-manifest.sql" "$container:$container_manifest" >/dev/null

coproc EXPORTER {{
  docker exec -i "$container" psql -X -U postgres -d "$database" -Atq -v ON_ERROR_STOP=1
}}
exporter_pid="$EXPORTER_PID"
exec 3>&"${{EXPORTER[1]}}"
exec 4<&"${{EXPORTER[0]}}"
printf '%s\n' 'begin transaction isolation level repeatable read read only;' 'select pg_export_snapshot();' >&3
IFS= read -r snapshot_id <&4
if [[ ! "$snapshot_id" =~ ^[0-9A-F]+-[0-9A-F]+-[0-9]+$ ]]; then
  echo "invalid exported snapshot id" >&2
  exit 1
fi
printf '%s\n' "select jsonb_build_object(
  'exported_snapshot_id', '$snapshot_id',
  'txid_snapshot', txid_current_snapshot()::text,
  'wal_lsn', pg_current_wal_lsn()::text,
  'system_identifier', (select system_identifier::text from pg_control_system()),
  'captured_at_utc', clock_timestamp()
)::text;" >&3
IFS= read -r snapshot_metadata <&4
if [[ -z "$snapshot_metadata" ]]; then
  echo "invalid source snapshot metadata" >&2
  exit 1
fi
printf '%s\n' "$snapshot_metadata" > "$remote_root/source-snapshot.json"

docker exec "$container" pg_dump \
  -U postgres -d "$database" \
  --format=custom --compress=9 --no-owner --no-acl \
  --snapshot="$snapshot_id" \
  --file="$container_dump"
docker exec "$container" psql \
  -X -U postgres -d "$database" -Atq -v ON_ERROR_STOP=1 \
  -v "snapshot_id=$snapshot_id" \
  -f "$container_manifest" \
  > "$remote_root/source-manifest.jsonl"

printf '%s\n' 'commit;' '\\q' >&3
exec 3>&-
while IFS= read -r _ <&4; do :; done
exec 4<&-
wait "$exporter_pid"

docker exec "$container" pg_restore --list "$container_dump" > "$remote_root/teleo-canonical.toc"
docker cp "$container:$container_dump" "$remote_root/teleo-canonical.dump" >/dev/null
sha256sum "$remote_root/teleo-canonical.dump" > "$remote_root/teleo-canonical.dump.sha256"
printf 'snapshot_exported=true\n' >> "$remote_root/source-context.txt"
systemctl show "$service" -p ActiveState -p SubState -p MainPID -p NRestarts -p ActiveEnterTimestamp --no-pager > "$remote_root/service-after.txt"
assert_service_healthy "$remote_root/service-after.txt"
cmp -s "$remote_root/service-before.txt" "$remote_root/service-after.txt"

tar -C "$remote_root" -cf - \
  teleo-canonical.dump \
  teleo-canonical.dump.sha256 \
  teleo-canonical.toc \
  source-manifest.jsonl \
  source-context.txt \
  source-snapshot.json \
  service-before.txt \
  service-after.txt
"""


def upload_manifest(ssh: list[str], target: str, remote_root: str, manifest: bytes) -> None:
    command = (
        f"umask 077; rm -rf {shlex.quote(remote_root)}; "
        f"mkdir -m 700 {shlex.quote(remote_root)}; "
        f"tee {shlex.quote(remote_root + '/postgres-parity-manifest.sql')} >/dev/null"
    )
    completed = subprocess.run(
        [*ssh, target, "--", command],
        input=manifest,
        capture_output=True,
        check=False,
    )
    if completed.returncode != 0:
        raise RuntimeError("manifest upload failed: " + completed.stderr.decode(errors="replace")[-1000:])


def cleanup_remote(ssh: list[str], target: str, remote_root: str) -> None:
    subprocess.run(
        [*ssh, target, "--", f"rm -rf {shlex.quote(remote_root)}"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
    )


def extract_capture(payload: bytes, output_dir: Path) -> None:
    seen: set[str] = set()
    with tarfile.open(fileobj=io.BytesIO(payload), mode="r:") as archive:
        for member in archive.getmembers():
            if not member.isfile() or member.name not in EXPECTED_FILES:
                raise RuntimeError(f"unexpected remote capture member: {member.name}")
            if member.name in seen:
                raise RuntimeError(f"duplicate remote capture member: {member.name}")
            source = archive.extractfile(member)
            if source is None:
                raise RuntimeError(f"missing remote capture payload: {member.name}")
            destination = output_dir / member.name
            descriptor = os.open(destination, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
            with os.fdopen(descriptor, "wb") as handle:
                handle.write(source.read())
            seen.add(member.name)
    if seen != EXPECTED_FILES:
        raise RuntimeError(f"incomplete remote capture: missing={sorted(EXPECTED_FILES - seen)}")


def toc_counts(path: Path) -> dict[str, int]:
    lines = [line for line in path.read_text(errors="replace").splitlines() if line and not line.startswith(";")]
    return {
        "entries": len(lines),
        "table_data": sum(" TABLE DATA " in line for line in lines),
        "constraints": sum(" CONSTRAINT " in line or " FK CONSTRAINT " in line for line in lines),
        "indexes": sum(" INDEX " in line for line in lines),
    }


def capture(args: argparse.Namespace) -> dict[str, Any]:
    output_dir = args.output_dir.resolve()
    if output_dir.exists():
        raise ValueError("--output-dir must not already exist")
    output_dir.parent.mkdir(parents=True, exist_ok=True)
    output_dir.mkdir(mode=0o700)
    remote_root = f"/tmp/teleo-canonical-snapshot-{args.run_id}"
    ssh = ssh_base(args.ssh_key.resolve())
    manifest_sql = args.manifest.resolve().read_bytes()
    manifest_sql_sha256 = sha256_bytes(manifest_sql)
    try:
        upload_manifest(ssh, args.ssh_target, remote_root, manifest_sql)
        script = remote_capture_script(
            run_id=args.run_id,
            container=args.container,
            database=args.database,
            service=args.service,
        )
        completed = subprocess.run(
            [*ssh, args.ssh_target, "--", "bash", "-s"],
            input=script.encode(),
            capture_output=True,
            check=False,
            timeout=args.timeout,
        )
        if completed.returncode != 0:
            raise RuntimeError("remote snapshot capture failed: " + completed.stderr.decode(errors="replace")[-2000:])
        extract_capture(completed.stdout, output_dir)
    finally:
        cleanup_remote(ssh, args.ssh_target, remote_root)

    dump_path = output_dir / "teleo-canonical.dump"
    expected_sha = (output_dir / "teleo-canonical.dump.sha256").read_text().split()[0]
    actual_sha = sha256(dump_path)
    if expected_sha != actual_sha:
        raise RuntimeError("captured dump SHA-256 mismatch")
    manifest = load_manifest(output_dir / "source-manifest.jsonl")
    snapshot = load_snapshot_metadata(output_dir / "source-snapshot.json")
    before = load_service_state(output_dir / "service-before.txt")
    after = load_service_state(output_dir / "service-after.txt")
    source = {
        "ssh_target": args.ssh_target,
        "container": args.container,
        "database": args.database,
        "service": args.service,
    }
    source_manifest_sha256 = sha256(output_dir / "source-manifest.jsonl")
    source_context_sha256 = sha256(output_dir / "source-context.txt")
    provenance_binding = build_provenance_binding(
        run_id=args.run_id,
        authorization_ref=args.authorization_ref,
        source=source,
        snapshot=snapshot,
        dump_sha256=actual_sha,
        manifest_sql_sha256=manifest_sql_sha256,
        source_manifest_sha256=source_manifest_sha256,
        source_context_sha256=source_context_sha256,
    )
    return {
        "artifact": "vps_canonical_postgres_exported_snapshot",
        "schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
        "receipt_version": 2,
        "generated_at_utc": datetime.now(UTC).isoformat(),
        "status": "pass",
        "run_id": args.run_id,
        "capture_authorization_ref": args.authorization_ref,
        "source": source,
        "snapshot_exported": True,
        "snapshot": snapshot,
        "dump": {
            "path": str(dump_path),
            "bytes": dump_path.stat().st_size,
            "sha256": actual_sha,
            "toc": toc_counts(output_dir / "teleo-canonical.toc"),
        },
        "manifest": {
            "path": str(output_dir / "source-manifest.jsonl"),
            "sha256": source_manifest_sha256,
            "manifest_sql_sha256": manifest_sql_sha256,
            "table_count": len(manifest["tables"]),
            "total_rows": sum(int(row["row_count"]) for row in manifest["tables"].values()),
            "application_roles": manifest["singleton"]["application_roles"]["items"],
        },
        "source_context": {
            "path": str(output_dir / "source-context.txt"),
            "sha256": source_context_sha256,
        },
        "source_service": {"before": before, "after": after, "unchanged": before == after},
        "provenance_binding": provenance_binding,
        "service_unchanged": before == after,
        "production_db_mutated": False,
        "telegram_send_attempted": False,
    }


def cleanup_failed_output(output_dir: Path, *, preexisting: bool) -> None:
    if not preexisting and output_dir.exists():
        shutil.rmtree(output_dir)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--execute", action="store_true")
    parser.add_argument("--ssh-target", required=True)
    parser.add_argument("--ssh-key", required=True, type=Path)
    parser.add_argument("--container", default="teleo-pg")
    parser.add_argument("--database", default="teleo")
    parser.add_argument("--service", default="leoclean-gateway.service")
    parser.add_argument("--manifest", default=Path("ops/postgres_parity_manifest.sql"), type=Path)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--authorization-ref", required=True)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--timeout", default=180, type=int)
    args = parser.parse_args(argv)
    for value, flag in ((args.container, "--container"), (args.database, "--database"), (args.service, "--service")):
        if not SAFE_NAME_RE.fullmatch(value):
            parser.error(f"{flag} contains unsafe characters")
    if not SAFE_SSH_TARGET_RE.fullmatch(args.ssh_target):
        parser.error("--ssh-target must be a host or user@host without options")
    if not SAFE_RUN_RE.fullmatch(args.run_id):
        parser.error("--run-id must be 8-64 safe characters")
    if not SAFE_AUTHORIZATION_REF_RE.fullmatch(args.authorization_ref):
        parser.error("--authorization-ref must be 8-256 safe metadata characters")
    if not args.ssh_key.is_file():
        parser.error("--ssh-key must be an existing file")
    if not args.manifest.is_file():
        parser.error("--manifest must be an existing file")
    if not args.execute:
        print(
            json.dumps(
                {
                    "artifact": "vps_canonical_postgres_exported_snapshot_plan",
                    "schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
                    "status": "dry_run",
                    "run_id": args.run_id,
                    "capture_authorization_ref": args.authorization_ref,
                    "source": f"{args.ssh_target}:{args.container}/{args.database}",
                    "manifest_sql_sha256": sha256(args.manifest.resolve()),
                    "output_dir": str(args.output_dir),
                    "mutates_production_db": False,
                },
                indent=2,
                sort_keys=True,
            )
        )
        raise SystemExit(0)
    return args


def main() -> int:
    args = parse_args()
    output_dir = args.output_dir.resolve()
    output_preexisted = output_dir.exists()
    try:
        payload = capture(args)
    except (OSError, ValueError, RuntimeError, subprocess.TimeoutExpired) as exc:
        cleanup_failed_output(output_dir, preexisting=output_preexisted)
        payload = {
            "artifact": "vps_canonical_postgres_exported_snapshot",
            "schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
            "generated_at_utc": datetime.now(UTC).isoformat(),
            "status": "fail",
            "run_id": args.run_id,
            "capture_authorization_ref": args.authorization_ref,
            "error": str(exc),
        }
        receipt = (
            output_dir.parent / f".{output_dir.name}-{args.run_id}-failure.json"
            if output_preexisted
            else output_dir / "receipt.json"
        )
        write_private_json(receipt, payload)
        print_private_receipt_summary(receipt, payload)
        return 1
    receipt = args.output_dir.resolve() / "receipt.json"
    write_private_json(receipt, payload)
    print_private_receipt_summary(receipt, payload)
    return 0


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