from __future__ import annotations

import argparse
import json
import shutil
import sqlite3
import tempfile
import zipfile
from collections.abc import Sequence
from contextlib import closing
from pathlib import Path

from tools.backup import sha256


def _safe_members(archive: zipfile.ZipFile) -> list[zipfile.ZipInfo]:
    members = archive.infolist()
    for member in members:
        path = Path(member.filename)
        if path.is_absolute() or ".." in path.parts:
            raise ValueError("backup contains an unsafe path")
    return members


def restore_backup(
    archive_path: Path,
    database: Path,
    *,
    private_skills: Path | None = None,
    environment_file: Path | None = None,
    encryption_key_file: Path | None = None,
    overwrite: bool = False,
) -> dict:
    archive_path = archive_path.resolve()
    database = database.resolve()
    destinations = [database]
    destinations.extend(path.resolve() for path in (environment_file, encryption_key_file) if path)
    if private_skills:
        destinations.append(private_skills.resolve())
    if not overwrite and any(path.exists() for path in destinations):
        raise FileExistsError(
            "restore destination exists; explicit overwrite confirmation is required"
        )
    with tempfile.TemporaryDirectory(prefix="xiaobai-restore-") as raw:
        staging = Path(raw)
        with zipfile.ZipFile(archive_path) as archive:
            _safe_members(archive)
            archive.extractall(staging)
        manifest = json.loads((staging / "manifest.json").read_text(encoding="utf-8"))
        if manifest.get("format") != 1 or not isinstance(manifest.get("files"), dict):
            raise ValueError("backup manifest is invalid")
        for name, metadata in manifest["files"].items():
            path = staging / name
            if not path.is_file() or sha256(path) != metadata.get("sha256"):
                raise ValueError(f"backup checksum mismatch: {name}")
        restored_db = staging / "database.sqlite3"
        with closing(
            sqlite3.connect(f"file:{restored_db.as_posix()}?mode=ro", uri=True)
        ) as connection:
            if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
                raise ValueError("restored database failed integrity check")
        database.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(restored_db, database)
        for destination, name in (
            (environment_file, "secrets/environment"),
            (encryption_key_file, "secrets/app-encryption.key"),
        ):
            source = staging / name
            if destination and source.is_file():
                destination.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(source, destination)
        skills_source = staging / "private-mentor-skills"
        if private_skills and skills_source.is_dir():
            private_skills.mkdir(parents=True, exist_ok=True)
            shutil.copytree(skills_source, private_skills, dirs_exist_ok=True)
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Validate and restore a Xiaobai backup")
    parser.add_argument("--archive", type=Path, required=True)
    parser.add_argument("--database", type=Path, required=True)
    parser.add_argument("--private-skills", type=Path)
    parser.add_argument("--environment-file", type=Path)
    parser.add_argument("--encryption-key-file", type=Path)
    parser.add_argument("--confirm-restore", action="store_true")
    parser.add_argument("--confirm-overwrite", action="store_true")
    return parser


def main(arguments: Sequence[str] | None = None) -> int:
    parsed = build_parser().parse_args(arguments)
    if not parsed.confirm_restore:
        raise SystemExit("restore requires --confirm-restore")
    manifest = restore_backup(
        parsed.archive,
        parsed.database,
        private_skills=parsed.private_skills,
        environment_file=parsed.environment_file,
        encryption_key_file=parsed.encryption_key_file,
        overwrite=parsed.confirm_overwrite,
    )
    print(
        json.dumps(
            {"restored": str(parsed.database.resolve()), "files": len(manifest["files"])},
            ensure_ascii=False,
        )
    )
    return 0


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