"""Back up AskMyDocs state to a Hugging Face Storage Bucket.

Creates a compressed archive of PostgreSQL data, Qdrant storage, and uploaded
documents, then uploads it to the specified HF Storage Bucket.  The Space's
``start.sh`` restores this archive on boot when ``RESTORE_ON_BOOT=true``.

Usage:
    python scripts/hf_state_backup.py --bucket username/askmydocs-data
    python scripts/hf_state_backup.py --bucket username/askmydocs-data --data-root /home/user
"""

from __future__ import annotations

import argparse
import subprocess
import tempfile
from pathlib import Path


def create_archive(data_root: Path, output: Path) -> bool:
    """Create a zstd-compressed tar of the state directories."""
    dirs_to_include = ["pgdata", "qdrant_storage"]
    uploads_dir = data_root / "app" / "data" / "uploads"
    if uploads_dir.exists():
        dirs_to_include.append("app/data/uploads")

    existing = [d for d in dirs_to_include if (data_root / d).exists()]
    if not existing:
        print("No state directories found to back up.")
        return False

    cmd = [
        "tar",
        "--zstd",
        "-cf",
        str(output),
        "-C",
        str(data_root),
        *existing,
    ]
    print(f"Creating archive: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"tar failed: {result.stderr}")
        return False

    size_mb = output.stat().st_size / (1024 * 1024)
    print(f"Archive created: {output} ({size_mb:.1f} MB)")
    return True


def upload_to_bucket(archive: Path, bucket: str, token: str | None = None) -> bool:
    """Upload the archive to an HF Storage Bucket.

    Storage Buckets are not repos: ``upload_file(repo_type="space_storage")``
    does not exist (supported repo types are model/dataset/space). The bucket
    API is ``batch_bucket_files(bucket_id=..., add=[(local, remote)])``.
    """
    try:
        from huggingface_hub import batch_bucket_files
    except ImportError:
        print("Error: huggingface_hub is required. Install with: pip install huggingface_hub")
        return False

    print(f"Uploading {archive.name} to bucket {bucket}...")
    batch_bucket_files(
        bucket_id=bucket,
        add=[(archive, "state.tar.zst")],
        token=token,
    )
    print(f"Uploaded to hf://buckets/{bucket}/state.tar.zst")
    return True


def main() -> int:
    parser = argparse.ArgumentParser(description="Back up AskMyDocs state to HF Storage Bucket.")
    parser.add_argument(
        "--bucket",
        required=True,
        help="HF Storage Bucket name (e.g. 'username/askmydocs-data')",
    )
    parser.add_argument(
        "--data-root",
        type=Path,
        default=Path("/home/user"),
        help="Data root directory (default: /home/user for HF Space)",
    )
    parser.add_argument(
        "--token",
        default=None,
        help="HF token for private bucket access (default: HF_TOKEN env var)",
    )
    args = parser.parse_args()

    import os
    token = args.token or os.environ.get("HF_TOKEN")

    with tempfile.NamedTemporaryFile(suffix=".tar.zst", delete=False) as tmp:
        tmp_path = Path(tmp.name)

    try:
        if not create_archive(args.data_root, tmp_path):
            return 1
        if not upload_to_bucket(tmp_path, args.bucket, token):
            return 1
    finally:
        tmp_path.unlink(missing_ok=True)

    print("State backup complete.")
    return 0


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