#!/usr/bin/env python3
import argparse
import json
import os
import shutil
import tarfile
import tempfile
from datetime import datetime, timezone
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

try:
    from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
except ImportError:
    from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError


DEFAULT_EXCLUDES = {
    ".cache",
    "__pycache__",
    "tmp",
    "temp",
}


def env(name: str, default: str | None = None) -> str | None:
    value = os.environ.get(name)
    return value if value not in (None, "") else default


def required_env(name: str) -> str:
    value = env(name)
    if not value:
        raise SystemExit(f"{name} is required")
    return value


def should_exclude(path: Path) -> bool:
    parts = set(path.parts)
    if parts & DEFAULT_EXCLUDES:
        return True
    name = path.name
    return name.endswith((".log", ".pid", ".pyc", ".tmp", ".lock"))


def state_has_content(root: Path) -> bool:
    if not root.exists():
        return False
    for child in root.rglob("*"):
        if child.is_file() and not should_exclude(child.relative_to(root)):
            return True
    return False


def state_stats(root: Path) -> tuple[int, int]:
    files = 0
    bytes_total = 0
    if not root.exists():
        return files, bytes_total
    for child in root.rglob("*"):
        if not child.is_file():
            continue
        rel = child.relative_to(root)
        if should_exclude(rel):
            continue
        files += 1
        try:
            bytes_total += child.stat().st_size
        except OSError:
            pass
    return files, bytes_total


def add_tree(tar: tarfile.TarFile, root: Path) -> None:
    for item in root.rglob("*"):
        rel = item.relative_to(root)
        if should_exclude(rel):
            continue
        tar.add(item, arcname=str(rel), recursive=False)


def validate_member(target: Path, member: tarfile.TarInfo) -> Path:
    destination = (target / member.name).resolve()
    target_resolved = target.resolve()
    if target_resolved != destination and target_resolved not in destination.parents:
        raise RuntimeError(f"Unsafe archive member: {member.name}")
    return destination


def safe_extract(archive: Path, target: Path) -> None:
    with tarfile.open(archive, "r:gz") as tar:
        for member in tar.getmembers():
            validate_member(target, member)
        tar.extractall(target)


def backup() -> None:
    repo = required_env("HERMES_BACKUP_REPO")
    token = required_env("HF_BACKUP_TOKEN")
    root = Path(env("HERMES_BACKUP_DIR", "/home/hermeswebui/.hermes")).resolve()
    filename = env("HERMES_BACKUP_FILENAME", "hermes-state.tar.gz")
    file_count, byte_count = state_stats(root)
    print(f"Hermes state backup source: {root} ({file_count} files, {byte_count} bytes).")

    if file_count == 0:
        print("Hermes state backup skipped: no meaningful local state yet.")
        return

    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    api = HfApi(token=token)

    with tempfile.TemporaryDirectory(prefix="hermes-backup-") as tmp:
        tmpdir = Path(tmp)
        archive = tmpdir / filename
        manifest = tmpdir / "manifest.json"

        with tarfile.open(archive, "w:gz") as tar:
            add_tree(tar, root)
        archive_size = archive.stat().st_size

        manifest.write_text(
            json.dumps(
                {
                    "created_at": timestamp,
                    "source": "Acrabohan/hermes-webui",
                    "path_in_repo": filename,
                    "file_count": file_count,
                    "source_bytes": byte_count,
                    "archive_bytes": archive_size,
                },
                indent=2,
            )
            + "\n",
            encoding="utf-8",
        )

        api.upload_file(
            path_or_fileobj=str(archive),
            path_in_repo=filename,
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message=f"Backup Hermes state {timestamp}",
        )
        api.upload_file(
            path_or_fileobj=str(manifest),
            path_in_repo="manifest.json",
            repo_id=repo,
            repo_type="dataset",
            token=token,
            commit_message=f"Update Hermes backup manifest {timestamp}",
        )

    print(f"Hermes state backup uploaded to dataset {repo} ({archive_size} bytes).")


def restore() -> None:
    repo = required_env("HERMES_BACKUP_REPO")
    token = required_env("HF_BACKUP_TOKEN")
    root = Path(env("HERMES_BACKUP_DIR", "/home/hermeswebui/.hermes")).resolve()
    filename = env("HERMES_BACKUP_FILENAME", "hermes-state.tar.gz")
    force = env("HERMES_FORCE_RESTORE", "0") == "1"

    if state_has_content(root) and not force:
        print("Hermes state restore skipped: local state already exists.")
        return

    root.mkdir(parents=True, exist_ok=True)

    try:
        archive_path = hf_hub_download(
            repo_id=repo,
            filename=filename,
            repo_type="dataset",
            token=token,
        )
    except EntryNotFoundError:
        print(f"Hermes state restore skipped: {filename} not found in {repo}.")
        return
    except HfHubHTTPError as exc:
        if getattr(exc.response, "status_code", None) == 404:
            print(f"Hermes state restore skipped: backup dataset or file not found.")
            return
        raise

    if force and state_has_content(root):
        stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
        backup_existing = root.parent / f"{root.name}.pre-restore-{stamp}"
        shutil.move(str(root), str(backup_existing))
        root.mkdir(parents=True, exist_ok=True)
        print(f"Existing Hermes state moved to {backup_existing}.")

    safe_extract(Path(archive_path), root)
    print(f"Hermes state restored from dataset {repo}.")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=("backup", "restore"))
    args = parser.parse_args()

    if args.command == "backup":
        backup()
    else:
        restore()


if __name__ == "__main__":
    main()