"""Restore a database backup.

SQLite restore copies a consistent snapshot into the target database file.
Stop the bot before restoring SQLite. PostgreSQL restores require pg_restore.

Examples:
    python scripts/restore_db.py backups/publisher.db --yes
    python scripts/restore_db.py backups/publisher.dump --yes --database-url postgresql+asyncpg://...
"""

from __future__ import annotations

import argparse
import os
import shutil
import sqlite3
import subprocess
import sys
from contextlib import closing
from pathlib import Path


def _sqlite_path(url: str) -> Path:
    prefix = "sqlite+aiosqlite:///"
    if not url.startswith(prefix):
        raise ValueError("unsupported SQLite URL")
    value = url[len(prefix) :]
    if value == ":memory:" or not value:
        raise ValueError("cannot restore into an in-memory SQLite database")
    return Path(value).expanduser()


def _postgres_url(url: str) -> str:
    return url.replace("postgresql+asyncpg://", "postgresql://", 1)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("backup", help="path to a SQLite .db file or PostgreSQL .dump file")
    parser.add_argument(
        "--database-url",
        default=os.environ.get("DATABASE_URL", "sqlite+aiosqlite:///./data/publisher.db"),
    )
    parser.add_argument("--yes", action="store_true", help="confirm destructive restore")
    args = parser.parse_args()

    if not args.yes:
        print("restore is destructive; pass --yes to continue", file=sys.stderr)
        return 2

    backup = Path(args.backup).expanduser()
    if not backup.exists():
        print(f"backup not found: {backup}", file=sys.stderr)
        return 2
    url = args.database_url.strip()

    if url.startswith("sqlite+aiosqlite:///"):
        target = _sqlite_path(url)
        target.parent.mkdir(parents=True, exist_ok=True)
        # Restore into a temporary file first so a failed copy cannot leave a
        # half-written live database.
        temporary = target.with_suffix(target.suffix + ".restore-tmp")
        if temporary.exists():
            temporary.unlink()
        with closing(
            sqlite3.connect(f"file:{backup.as_posix()}?mode=ro", uri=True)
        ) as src, closing(sqlite3.connect(temporary)) as dst:
            src.backup(dst)
        os.replace(temporary, target)
        print(f"SQLite database restored to {target}")
        return 0

    if url.startswith("postgresql+asyncpg://"):
        pg_restore = shutil.which("pg_restore")
        if pg_restore is None:
            print("pg_restore was not found on PATH", file=sys.stderr)
            return 2
        subprocess.run(
            [
                pg_restore,
                "--clean",
                "--if-exists",
                "--no-owner",
                "--no-privileges",
                f"--dbname={_postgres_url(url)}",
                str(backup),
            ],
            check=True,
        )
        print("PostgreSQL database restored")
        return 0

    print("DATABASE_URL must be sqlite+aiosqlite:/// or postgresql+asyncpg://", file=sys.stderr)
    return 2


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