#!/usr/bin/env python3
"""Atomic Hugging Face Dataset persistence for HermesFace state.

The public helpers retain the legacy single-file API, while the production
sync path commits one validated secure-backup archive plus metadata in a single
Dataset commit.
"""
from __future__ import annotations

import hashlib
import json
import logging
import os
import shutil
import sys
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any, Callable, Dict, List, Optional, TypeVar

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError

T = TypeVar("T")
_RETRYABLE_STATUS = {408, 425, 429, 500, 502, 503, 504}

logging.basicConfig(
    level=logging.INFO,
    format='{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "atomic-save", "message": "%(message)s"}',
)
logger = logging.getLogger(__name__)


def _is_retryable_hf_error(exc: Exception) -> bool:
    """True for transient Hugging Face Hub / network failures."""
    if isinstance(exc, HfHubHTTPError):
        response = getattr(exc, "response", None)
        status = getattr(response, "status_code", None)
        if status in _RETRYABLE_STATUS:
            return True
    message = str(exc).lower()
    return any(token in message for token in ("502", "503", "504", "bad gateway", "gateway timeout"))


def _safe_dataset_path(value: str) -> str:
    pure = PurePosixPath(str(value).replace("\\", "/"))
    if pure.is_absolute() or not pure.parts or any(part in {"", ".", ".."} for part in pure.parts):
        raise ValueError(f"unsafe dataset path: {value!r}")
    return pure.as_posix()


def validate_backup_source_paths(source_paths: List[str], state_dir: Optional[str] = None) -> List[Path]:
    """Resolve and validate explicitly allowlisted backup source files."""
    try:
        from backup_allowlist import BackupAllowlistError, is_path_allowed, scan_file_for_secrets
    except ImportError:
        from scripts.backup_allowlist import (  # type: ignore
            BackupAllowlistError,
            is_path_allowed,
            scan_file_for_secrets,
        )

    root = Path(state_dir or os.environ.get("HERMES_HOME", "/opt/data")).resolve()
    validated: List[Path] = []
    for raw in source_paths:
        unresolved = Path(raw)
        if unresolved.is_symlink():
            raise BackupAllowlistError(f"Refusing symlink backup source: {unresolved.name}")
        path = unresolved.resolve()
        try:
            relative = path.relative_to(root)
        except ValueError as exc:
            raise BackupAllowlistError(f"Refusing backup outside HERMES_HOME: {path.name}") from exc
        if not path.is_file():
            raise BackupAllowlistError(f"Refusing non-file backup source: {relative.as_posix()}")
        if not is_path_allowed(relative.as_posix()):
            raise BackupAllowlistError(f"Refusing non-allowlisted backup source: {relative.as_posix()}")
        reason = scan_file_for_secrets(path)
        if reason:
            raise BackupAllowlistError(
                f"Refusing backup: secret-like content in {relative.as_posix()} ({reason})"
            )
        validated.append(path)
    return validated


class AtomicDatasetSaver:
    def __init__(
        self,
        repo_id: str,
        dataset_path: str = "state",
        *,
        api: Optional[HfApi] = None,
        token: Optional[str] = None,
    ):
        self.repo_id = repo_id
        self.dataset_path = _safe_dataset_path(dataset_path)
        self.token = token
        self.api = api or HfApi(token=token)
        self.max_retries = 3
        self.base_delay = 1.0
        self.max_backups = 3
        logger.info(f"init repo_id={repo_id} dataset_path={self.dataset_path}")

    def _call_with_retry(self, label: str, fn: Callable[[], T]) -> T:
        """Retry transient Hub API failures with bounded exponential backoff."""
        attempt = 0
        while True:
            try:
                return fn()
            except Exception as exc:
                attempt += 1
                if attempt > self.max_retries or not _is_retryable_hf_error(exc):
                    raise
                delay = min(30.0, self.base_delay * (2 ** (attempt - 1)))
                logger.warning(
                    f"{label}_retry attempt={attempt}/{self.max_retries} "
                    f"delay_s={delay:.1f} error={type(exc).__name__}"
                )
                time.sleep(delay)

    def calculate_checksum(self, file_path: Path) -> str:
        digest = hashlib.sha256()
        with file_path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        return digest.hexdigest()

    def _repo_info(self):
        try:
            return self._call_with_retry(
                "repo_info",
                lambda: self.api.repo_info(repo_id=self.repo_id, repo_type="dataset"),
            )
        except RepositoryNotFoundError:
            return None

    def create_backup(self, current_commit: Optional[str]) -> Optional[str]:
        """Copy the current authoritative state files to a timestamped remote backup."""
        if not current_commit:
            return None
        try:
            files = self._call_with_retry(
                "list_repo_files",
                lambda: self.api.list_repo_files(
                    repo_id=self.repo_id, repo_type="dataset", revision=current_commit
                ),
            )
            prefix = f"{self.dataset_path}/"
            state_files = [path for path in files if path.startswith(prefix)]
            if not state_files:
                return None

            timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
            backup_path = f"backups/state_{timestamp}"
            logger.info(f"creating_backup path={backup_path} files={len(state_files)}")

            with tempfile.TemporaryDirectory(prefix="hermes-hf-backup-") as tmpdir:
                root = Path(tmpdir)
                operations = []
                for remote_path in state_files:
                    local = self._call_with_retry(
                        "hf_hub_download",
                        lambda remote_path=remote_path: hf_hub_download(
                            repo_id=self.repo_id,
                            repo_type="dataset",
                            filename=remote_path,
                            revision=current_commit,
                            token=self.token,
                        ),
                    )
                    relative = PurePosixPath(remote_path).relative_to(PurePosixPath(self.dataset_path))
                    destination = root.joinpath(*relative.parts)
                    destination.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(local, destination)
                    operations.append(
                        CommitOperationAdd(
                            path_in_repo=f"{backup_path}/{relative.as_posix()}",
                            path_or_fileobj=str(destination),
                        )
                    )

                if not operations:
                    return None
                info = self._call_with_retry(
                    "create_backup_commit",
                    lambda: self.api.create_commit(
                        repo_id=self.repo_id,
                        repo_type="dataset",
                        operations=operations,
                        commit_message=f"Backup state before update - {timestamp}",
                        parent_commit=current_commit,
                    ),
                )
                logger.info(f"backup_created commit={info.oid}")
                return info.oid
        except Exception as exc:
            # A failed remote backup must prevent replacing the current remote state.
            logger.error(f"backup_failed error={exc}")
            raise

    def save_secure_archive_atomic(
        self,
        archive_path: Path,
        state_data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Validate and atomically commit one secure backup archive and metadata."""
        try:
            from secure_backup import extract_validated
        except ImportError:
            from scripts.secure_backup import extract_validated  # type: ignore

        archive = Path(archive_path)
        if archive.is_symlink() or not archive.is_file():
            raise ValueError("secure backup archive must be a regular file")

        # Validate member paths, member types, manifest, hashes, and SQLite before
        # any Hugging Face API call is permitted.
        with tempfile.TemporaryDirectory(prefix="hermes-archive-validate-") as tmpdir:
            manifest = extract_validated(archive, Path(tmpdir) / "validated")

        archive_sha256 = self.calculate_checksum(archive)
        operation_id = f"save_{int(time.time())}"
        repo_info = self._repo_info()
        current_commit = getattr(repo_info, "sha", None) if repo_info else None
        backup_commit = self.create_backup(current_commit)
        parent_commit = backup_commit or current_commit

        metadata = {
            "format": "hermes-secure-backup",
            "format_version": 1,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "operation_id": operation_id,
            "archive": "secure_state.tar.gz",
            "archive_sha256": archive_sha256,
            "manifest_sha256": hashlib.sha256(
                json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8")
            ).hexdigest(),
            "manifest": manifest,
            "backup_commit": backup_commit,
            "state_data": state_data or {},
        }

        with tempfile.TemporaryDirectory(prefix="hermes-hf-save-") as tmpdir:
            metadata_path = Path(tmpdir) / "metadata.json"
            metadata_path.write_text(
                json.dumps(metadata, sort_keys=True, separators=(",", ":")),
                encoding="utf-8",
            )
            operations = [
                CommitOperationAdd(
                    path_in_repo=f"{self.dataset_path}/secure_state.tar.gz",
                    path_or_fileobj=str(archive),
                ),
                CommitOperationAdd(
                    path_in_repo=f"{self.dataset_path}/metadata.json",
                    path_or_fileobj=str(metadata_path),
                ),
            ]
            info = self._call_with_retry(
                "create_commit",
                lambda: self.api.create_commit(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    operations=operations,
                    commit_message=f"Atomic secure state update - {operation_id}",
                    parent_commit=parent_commit,
                ),
            )

        result = {
            "success": True,
            "operation_id": operation_id,
            "commit_id": info.oid,
            "backup_commit": backup_commit,
            "timestamp": metadata["timestamp"],
            "archive_sha256": archive_sha256,
            "files_count": len(manifest.get("files", [])),
        }
        logger.info(f"atomic_secure_save_completed {result}")
        return result

    def save_state_atomic(
        self, state_data: Dict[str, Any], source_paths: List[str]
    ) -> Dict[str, Any]:
        """Legacy allowlisted individual-file commit retained for compatibility."""
        operation_id = f"save_{int(time.time())}"
        logger.info(f"starting_atomic_save op={operation_id} sources={source_paths}")
        validated_sources = validate_backup_source_paths(source_paths)

        repo_info = self._repo_info()
        current_commit = getattr(repo_info, "sha", None) if repo_info else None
        backup_commit = self.create_backup(current_commit)
        parent_commit = backup_commit or current_commit

        with tempfile.TemporaryDirectory(prefix="hermes-state-save-") as tmpdir:
            state_dir = Path(tmpdir) / "state"
            state_dir.mkdir(parents=True, exist_ok=True)
            metadata = {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "operation_id": operation_id,
                "checksum": hashlib.sha256(
                    json.dumps(state_data, sort_keys=True).encode("utf-8")
                ).hexdigest(),
                "backup_commit": backup_commit,
                "state_data": state_data,
            }
            metadata_path = state_dir / "metadata.json"
            metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
            operations = [
                CommitOperationAdd(
                    path_in_repo=f"{self.dataset_path}/metadata.json",
                    path_or_fileobj=str(metadata_path),
                )
            ]
            for source in validated_sources:
                destination = state_dir / source.name
                shutil.copy2(source, destination)
                operations.append(
                    CommitOperationAdd(
                        path_in_repo=f"{self.dataset_path}/{source.name}",
                        path_or_fileobj=str(destination),
                    )
                )
            info = self._call_with_retry(
                "create_commit",
                lambda: self.api.create_commit(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    operations=operations,
                    commit_message=f"Atomic state update - {operation_id}",
                    parent_commit=parent_commit,
                ),
            )

        return {
            "success": True,
            "operation_id": operation_id,
            "commit_id": info.oid,
            "backup_commit": backup_commit,
            "timestamp": metadata["timestamp"],
            "files_count": len(validated_sources),
        }


def main() -> None:
    if len(sys.argv) < 3:
        print(json.dumps({
            "error": "Usage: python save_to_dataset_atomic.py <repo_id> <source_path1> [source_path2...]",
            "status": "error",
        }, indent=2))
        sys.exit(1)

    repo_id = sys.argv[1]
    source_paths = sys.argv[2:]
    for path in source_paths:
        if not os.path.exists(path):
            print(json.dumps({"error": f"Source path does not exist: {path}", "status": "error"}, indent=2))
            sys.exit(1)

    state_data = {
        "environment": "production",
        "version": "1.0.0",
        "platform": "huggingface-spaces",
        "app": "hermesface",
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    try:
        result = AtomicDatasetSaver(repo_id).save_state_atomic(state_data, source_paths)
        print(json.dumps(result, indent=2))
    except Exception as exc:
        print(json.dumps({"error": str(exc), "status": "error"}, indent=2))
        sys.exit(1)


if __name__ == "__main__":
    main()