from __future__ import annotations

import argparse
import gzip
import hashlib
import shutil
from pathlib import Path


DEFAULT_SNAPSHOT = Path(__file__).resolve().parents[1] / "snapshots" / "wehago_account_ledger_snapshot_20260605.sqlite3.gz"
DEFAULT_OUTPUT = Path(__file__).resolve().parents[1] / "data" / "wehago_account_ledger.sqlite3"
EXPECTED_SHA256 = "dd1e8792d277427f1c95aa95f8eeefe92b208adc51168161ec3ec5b16f0b526f"


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 main() -> None:
    parser = argparse.ArgumentParser(description="Restore the shared WEHAGO account ledger SQLite snapshot.")
    parser.add_argument("--snapshot", type=Path, default=DEFAULT_SNAPSHOT, help="Path to the .sqlite3.gz snapshot.")
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="Output SQLite DB path.")
    parser.add_argument("--skip-checksum", action="store_true", help="Skip snapshot checksum verification.")
    args = parser.parse_args()

    snapshot = args.snapshot.expanduser().resolve()
    output = args.output.expanduser().resolve()
    if not snapshot.exists():
        raise SystemExit(f"Snapshot not found: {snapshot}")
    if not args.skip_checksum:
        actual = sha256(snapshot)
        if actual != EXPECTED_SHA256:
            raise SystemExit(f"Checksum mismatch: {actual} != {EXPECTED_SHA256}")

    output.parent.mkdir(parents=True, exist_ok=True)
    temp_output = output.with_suffix(output.suffix + ".tmp")
    with gzip.open(snapshot, "rb") as source, temp_output.open("wb") as dest:
        shutil.copyfileobj(source, dest)
    temp_output.replace(output)
    print(f"Restored: {output}")


if __name__ == "__main__":
    main()
