"""cortex/storage/backup.py — Backup and restore logic for Cortex SQLite database.

Provides:
  - backup_database(): create a timestamped tar.gz backup with SHA-256 checksum
  - restore_database(): verify checksum, restore from tar.gz, run PRAGMA integrity_check
  - list_backups(): enumerate existing backups in a directory
  - prune_backups(): remove old backups keeping N most recent
"""
from __future__ import annotations

import hashlib
import logging
import os
import pathlib
import re
import shutil
import sqlite3
import sys
import tarfile
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

CHECKSUM_SUFFIX = ".sha256"
BACKUP_PATTERN = re.compile(r"^.+_\d{8}T\d{6}Z\.tar\.gz$")


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _sha256_file(path: Path) -> str:
    """Return the hex SHA-256 digest of a file."""
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def _write_checksum(archive_path: Path) -> Path:
    """Write a .sha256 sidecar file for *archive_path* and return its path."""
    digest = _sha256_file(archive_path)
    checksum_path = archive_path.with_suffix(".tar.gz" + CHECKSUM_SUFFIX)
    # e.g. cortex_20260101T120000Z.tar.gz.sha256
    checksum_path.write_text(f"{digest}  {archive_path.name}\n", encoding="utf-8")
    logger.debug("Wrote checksum %s → %s", digest, checksum_path)
    return checksum_path


def _verify_checksum(archive_path: Path) -> tuple[bool, str]:
    """
    Verify the SHA-256 checksum of *archive_path* against its sidecar.

    Returns (ok: bool, message: str).
    """
    checksum_path = archive_path.with_suffix(".tar.gz" + CHECKSUM_SUFFIX)
    if not checksum_path.exists():
        return False, f"Checksum file not found: {checksum_path}"

    stored_line = checksum_path.read_text(encoding="utf-8").strip()
    parts = stored_line.split(None, 1)
    if len(parts) < 1:
        return False, "Checksum file is malformed"

    stored_digest = parts[0]
    actual_digest = _sha256_file(archive_path)

    if stored_digest == actual_digest:
        return True, "Checksum OK"
    return False, f"Checksum mismatch: expected {stored_digest}, got {actual_digest}"


def _safe_extractall(tar: tarfile.TarFile, dest: Path) -> None:
    """Extract *tar* to *dest* with path-traversal protection (CVE-2007-4559).

    On Python ≥ 3.12 we use the built-in ``filter='data'`` parameter.
    On older versions we manually validate every member name.
    """
    if sys.version_info >= (3, 12):
        tar.extractall(path=dest, filter="data")
    else:
        dest_resolved = pathlib.Path(dest).resolve()
        for member in tar.getmembers():
            if member.name.startswith("..") or member.name.startswith("/"):
                raise ValueError(f"Unsafe path in archive: {member.name}")
            resolved = dest_resolved.joinpath(member.name).resolve()
            if not str(resolved).startswith(str(dest_resolved)):
                raise ValueError(f"Path traversal attempt: {member.name}")
        tar.extractall(path=dest)  # noqa: S202 — validated above


class WALCheckpointError(RuntimeError):
    """Raised when a WAL checkpoint fails and the caller requires success."""


def _wal_checkpoint(db_path: Path, timeout: float = 10) -> None:
    """Issue a WAL checkpoint to flush WAL frames to the main database file.

    Uses PRAGMA wal_checkpoint(TRUNCATE) which requires exclusive access.
    The *timeout* parameter (seconds) limits how long we wait for the database
    lock — this prevents blocking writers indefinitely.  For best results,
    run backups during low-traffic windows.

    Raises:
        WALCheckpointError: If the checkpoint fails for any reason.
    """
    try:
        conn = sqlite3.connect(str(db_path), timeout=timeout)
        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
        conn.close()
        logger.debug("WAL checkpoint complete on %s", db_path)
    except sqlite3.Error as exc:
        logger.error("WAL checkpoint failed: %s", exc)
        raise WALCheckpointError(f"WAL checkpoint failed: {exc}") from exc


def _integrity_check(db_path: Path) -> tuple[bool, list[str]]:
    """
    Run PRAGMA integrity_check on the database at *db_path*.

    Returns (ok: bool, messages: list[str]).
    """
    try:
        conn = sqlite3.connect(str(db_path), timeout=30)
        rows = conn.execute("PRAGMA integrity_check").fetchall()
        conn.close()
        messages = [r[0] for r in rows]
        ok = messages == ["ok"]
        return ok, messages
    except sqlite3.Error as exc:
        return False, [str(exc)]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def backup_database(
    db_path: Path,
    output_dir: Path,
    name_prefix: str = "cortex",
    include_config: Optional[Path] = None,
    keep: Optional[int] = None,
    require_checkpoint: bool = True,
) -> tuple[Path, Path]:
    """
    Create a compressed, checksummed backup of the Cortex SQLite database.

    Steps:
      1. WAL checkpoint (flush WAL → main DB file for consistency)
      2. Copy DB (and optional WAL/SHM files) into a temp dir
      3. Optionally include a config file
      4. Pack into a timestamped tar.gz
      5. Write SHA-256 sidecar
      6. Prune old backups if --keep N supplied

    Args:
        db_path:        Path to the live SQLite database.
        output_dir:     Directory where the backup archive will be written.
        name_prefix:    Prefix for the archive filename (default: "cortex").
        include_config: Optional path to cortex.toml to bundle into the archive.
        keep:           If set, keep only the N most recent backups (prune older ones).
        require_checkpoint: If True (default), abort the backup when the WAL
                        checkpoint fails.  Set to False to proceed anyway (the
                        backup will still be consistent but may not include the
                        very latest WAL frames).

    Returns:
        (archive_path, checksum_path) — both Path objects.

    Raises:
        WALCheckpointError: If ``require_checkpoint`` is True and the WAL
            checkpoint fails.
    """
    db_path = Path(db_path).resolve()
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if not db_path.exists():
        raise FileNotFoundError(f"Database not found: {db_path}")

    # Step 1: WAL checkpoint
    try:
        _wal_checkpoint(db_path)
    except WALCheckpointError:
        if require_checkpoint:
            raise
        logger.warning("WAL checkpoint failed but require_checkpoint=False — proceeding")

    # Step 2: Stage files in a temp dir
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
    archive_name = f"{name_prefix}_{timestamp}.tar.gz"
    archive_path = output_dir / archive_name

    with tempfile.TemporaryDirectory(prefix="cortex-backup-") as tmp:
        tmp_path = Path(tmp)
        staging = tmp_path / f"{name_prefix}_{timestamp}"
        staging.mkdir()

        # Copy main DB
        shutil.copy2(db_path, staging / db_path.name)

        # Copy WAL/SHM if present (best-effort — after checkpoint they should be tiny/empty)
        for ext in ("-wal", "-shm"):
            companion = db_path.with_name(db_path.name + ext)
            if companion.exists():
                shutil.copy2(companion, staging / companion.name)

        # Step 3: Optional config
        if include_config is not None:
            cfg_path = Path(include_config).resolve()
            if cfg_path.exists():
                shutil.copy2(cfg_path, staging / cfg_path.name)
            else:
                logger.warning("Config file not found, skipping: %s", cfg_path)

        # Step 4: Create tar.gz in temp dir first, then move to output_dir
        # on success.  This prevents partial archives on disk if tar creation
        # fails (e.g. disk full).
        tmp_archive = tmp_path / archive_name
        with tarfile.open(tmp_archive, "w:gz") as tar:
            tar.add(staging, arcname=staging.name)

        shutil.move(str(tmp_archive), str(archive_path))

    logger.info("Backup archive created: %s", archive_path)

    # Step 5: Write checksum
    checksum_path = _write_checksum(archive_path)
    logger.info("Checksum written: %s", checksum_path)

    # Step 6: Prune
    if keep is not None and keep > 0:
        pruned = prune_backups(output_dir, name_prefix=name_prefix, keep=keep)
        if pruned:
            logger.info("Pruned %d old backup(s)", len(pruned))

    return archive_path, checksum_path


def restore_database(
    archive_path: Path,
    db_path: Path,
    force: bool = False,
    skip_checksum: bool = False,
) -> None:
    """
    Restore a Cortex database from a backup archive.

    **Important:** The target database must be quiesced (service stopped) before
    calling this function.  If WAL or SHM files exist at the target path, the
    database may still be in use and the restore will abort unless ``force=True``.

    Steps:
      1. Check for active WAL/SHM files (abort if present and not forced)
      2. Verify SHA-256 checksum (unless --force / skip_checksum)
      3. Extract archive to a temp directory
      4. Locate the .db file inside the archive
      5. Move existing DB aside (db_path.bak)
      6. Copy restored DB into place
      7. Run PRAGMA integrity_check
      8. On failure, roll back to the .bak file

    Args:
        archive_path:   Path to the .tar.gz backup archive.
        db_path:        Destination database path to restore into.
        force:          If True, proceed even when WAL/SHM files exist at the
                        target path or when checksum verification fails.
                        (The checksum is still *checked*, but a mismatch becomes
                        a warning instead of an error.)  Has no effect when
                        skip_checksum is also True.
        skip_checksum:  If True, bypass checksum verification entirely (not
                        recommended).

    Raises:
        FileNotFoundError: Archive not found.
        RuntimeError: WAL/SHM files exist at target and force is False (database
            may be in use).
        ValueError: Checksum mismatch (unless skip_checksum=True), or unsafe
            archive member paths (CVE-2007-4559).
        RuntimeError: Integrity check failed after restore.
    """
    archive_path = Path(archive_path).resolve()
    db_path = Path(db_path).resolve()

    if not archive_path.exists():
        raise FileNotFoundError(f"Archive not found: {archive_path}")

    # Step 0: Check for active WAL/SHM files — indicates DB may be in use
    wal_path = db_path.with_name(db_path.name + "-wal")
    shm_path = db_path.with_name(db_path.name + "-shm")
    active_files = [p for p in (wal_path, shm_path) if p.exists()]
    if active_files:
        names = ", ".join(str(p) for p in active_files)
        if not force:
            raise RuntimeError(
                f"Database may be in use — WAL/SHM files exist: {names}. "
                f"Stop the service before restoring, or pass force=True to proceed anyway."
            )
        logger.warning(
            "WAL/SHM files exist at target (%s) — proceeding because force=True", names
        )

    # Step 1: Checksum verification
    if not skip_checksum:
        ok, msg = _verify_checksum(archive_path)
        if not ok:
            if not force:
                raise ValueError(f"Checksum verification failed: {msg}")
            else:
                logger.warning("Checksum verification failed (proceeding with --force): %s", msg)
        else:
            logger.info("Checksum verified OK")

    # Step 2: Extract to temp (with path-traversal protection — CVE-2007-4559)
    with tempfile.TemporaryDirectory(prefix="cortex-restore-") as tmp:
        tmp_path = Path(tmp)
        with tarfile.open(archive_path, "r:gz") as tar:
            _safe_extractall(tar, tmp_path)

        # Step 3: Find .db file inside extracted directory
        db_files = list(tmp_path.rglob("*.db"))
        if not db_files:
            raise FileNotFoundError("No .db file found inside backup archive")
        if len(db_files) > 1:
            logger.warning("Multiple .db files found in archive, using: %s", db_files[0])
        restored_db = db_files[0]

        # Step 4: Back up current DB (atomic)
        backup_existing: Optional[Path] = None
        if db_path.exists():
            backup_existing = db_path.with_suffix(".db.bak")
            tmp_bak = db_path.with_suffix(".db.bak.tmp")
            shutil.copy2(db_path, tmp_bak)
            os.replace(tmp_bak, backup_existing)
            logger.info("Existing DB backed up to: %s", backup_existing)

        try:
            # Step 5: Copy restored DB into place (atomic via temp + os.replace)
            db_path.parent.mkdir(parents=True, exist_ok=True)
            tmp_restore = db_path.with_suffix(".db.tmp")
            shutil.copy2(restored_db, tmp_restore)
            os.replace(tmp_restore, db_path)
            logger.info("Restored DB to: %s", db_path)

            # Also restore WAL/SHM if present in archive (but remove stale ones first)
            for ext in ("-wal", "-shm"):
                stale = db_path.with_name(db_path.name + ext)
                if stale.exists():
                    stale.unlink()
                archived_companion = restored_db.with_name(restored_db.name + ext)
                if archived_companion.exists():
                    shutil.copy2(archived_companion, db_path.with_name(db_path.name + ext))

            # Step 6: Integrity check
            ok, messages = _integrity_check(db_path)
            if not ok:
                raise RuntimeError(
                    f"Integrity check failed after restore: {'; '.join(messages)}"
                )
            logger.info("Integrity check passed")

        except Exception:
            # Step 7: Roll back (atomic)
            if backup_existing and backup_existing.exists():
                tmp_rollback = db_path.with_suffix(".db.rollback.tmp")
                shutil.copy2(backup_existing, tmp_rollback)
                os.replace(tmp_rollback, db_path)
                logger.error("Restore failed — rolled back to previous database")
            raise


def list_backups(output_dir: Path, name_prefix: str = "cortex") -> list[Path]:
    """
    Return a list of backup archives in *output_dir* matching *name_prefix*,
    sorted by modification time (oldest first).
    """
    output_dir = Path(output_dir)
    if not output_dir.exists():
        return []
    archives = [
        p for p in output_dir.iterdir()
        if p.name.startswith(name_prefix) and p.name.endswith(".tar.gz")
    ]
    archives.sort(key=lambda p: p.stat().st_mtime)
    return archives


def prune_backups(output_dir: Path, name_prefix: str = "cortex", keep: int = 5) -> list[Path]:
    """
    Remove old backup archives (and their checksum sidecars) from *output_dir*,
    keeping the *keep* most recent ones.

    Returns a list of removed archive paths.
    """
    archives = list_backups(output_dir, name_prefix=name_prefix)
    to_remove = archives[:-keep] if keep > 0 else archives

    removed: list[Path] = []
    for archive in to_remove:
        try:
            archive.unlink()
            removed.append(archive)
            # Also remove sidecar checksum if present
            sidecar = archive.with_suffix(".tar.gz" + CHECKSUM_SUFFIX)
            if sidecar.exists():
                sidecar.unlink()
        except OSError as exc:
            logger.warning("Could not remove backup %s: %s", archive, exc)

    return removed
