deploy.py

#!/usr/bin/env python3
"""
deploy.py — Persistent data deployment for Hugging Face Spaces.

Usage:
    python deploy.py

Flow:
    1. Back up existing data from the deployed Space -> persistence/
    2. Commit and push local changes to the Space repository
    3. Wait for the Space to redeploy (poll /health)
    4. Restore the backed-up data to the new deployment
    5. Generate a comprehensive audit report
"""

import argparse
import io
import json
import os
import re
import subprocess
import sys
import tarfile
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

# ─────────────────────────────────────────────────────────────
# Console encoding detection
# ─────────────────────────────────────────────────────────────

_USE_UNICODE = True
try:
    "\u2713".encode(sys.stdout.encoding or "utf-8")
except (UnicodeEncodeError, UnicodeDecodeError, LookupError):
    _USE_UNICODE = False

# ─────────────────────────────────────────────────────────────
# ANSI colours & helpers
# ─────────────────────────────────────────────────────────────

GREEN = "\033[92m"
RED = "\033[91m"
YELLOW = "\033[93m"
CYAN = "\033[96m"
BOLD = "\033[1m"
RESET = "\033[0m"


def _c(colour: str, text: str) -> str:
    return f"{colour}{text}{RESET}"


def _ok(text: str) -> str:
    return _c(GREEN, text)


def _fail(text: str) -> str:
    return _c(RED, text)


def _warn(text: str) -> str:
    return _c(YELLOW, text)


def _info(text: str) -> str:
    return _c(CYAN, text)


def _bold(text: str) -> str:
    return _c(BOLD, text)


# Icon constants (must be defined outside f-string expressions for Python < 3.12)
# Unicode vs ASCII fallbacks for Windows console compatibility
if _USE_UNICODE:
    _OK_ICON = _ok("\u2713")
    _FAIL_ICON = _fail("\u2717")
    _WARN_ICON = _warn("\u26A0")
    _INFO_ICON = _info("\u25B8")
    _BULLET = "\u2022"
    _DASH = "\u2014"
    _TABLE_TL = "\u2554"
    _TABLE_TR = "\u2557"
    _TABLE_BL = "\u255A"
    _TABLE_BR = "\u255D"
    _TABLE_H = "\u2550"
    _TABLE_V = "\u2551"
    _TABLE_TM = "\u2560"
    _TABLE_BM = "\u255A"
    _TABLE_ML = "\u251C"
    _TABLE_MR = "\u2524"
    _TABLE_MM = "\u253C"
    _TABLE_MH = "\u2500"
    _TABLE_MV = "\u2502"
    _TABLE_TML = "\u255F"
    _TABLE_TMR = "\u257E"
else:
    _OK_ICON = _ok("v")
    _FAIL_ICON = _fail("x")
    _WARN_ICON = _warn("!")
    _INFO_ICON = _info(">")
    _BULLET = "*"
    _DASH = "-"
    _TABLE_TL = "+"
    _TABLE_TR = "+"
    _TABLE_BL = "+"
    _TABLE_BR = "+"
    _TABLE_H = "="
    _TABLE_V = "|"
    _TABLE_TM = "+"
    _TABLE_BM = "+"
    _TABLE_ML = "+"
    _TABLE_MR = "+"
    _TABLE_MM = "+"
    _TABLE_MH = "-"
    _TABLE_MV = "|"
    _TABLE_TML = "+"
    _TABLE_TMR = "-"


def _indicator(status: str) -> str:
    mapping = {
        "passed": _OK_ICON,
        "failed": _FAIL_ICON,
        "skipped": _WARN_ICON,
        "running": _INFO_ICON,
    }
    return mapping.get(status, _DASH)


def _timestamp() -> str:
    return datetime.now(timezone.utc).strftime("%H:%M:%S")


def _human_size(n_bytes: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n_bytes < 1024:
            return f"{n_bytes:.2f} {unit}"
        n_bytes /= 1024
    return f"{n_bytes:.2f} TB"


def _human_duration(seconds: float) -> str:
    if seconds < 60:
        return f"{seconds:.1f}s"
    mins = int(seconds // 60)
    secs = int(seconds % 60)
    return f"{mins}m {secs}s"


# ─────────────────────────────────────────────────────────────
# Environment loading
# ─────────────────────────────────────────────────────────────


def load_env(env_path: str = ".env") -> dict[str, str]:
    env: dict[str, str] = {}
    p = Path(env_path)
    if not p.is_file():
        return env
    for line in p.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        m = re.match(r"^([A-Za-z_][A-Za-z_0-9]*)\s*=\s*(.*?)\s*$", line)
        if not m:
            continue
        key = m.group(1)
        val = m.group(2)
        if val.startswith('"') and val.endswith('"'):
            val = val[1:-1]
        elif val.startswith("'") and val.endswith("'"):
            val = val[1:-1]
        env[key] = val
    return env


# ─────────────────────────────────────────────────────────────
# HTTP helpers
# ─────────────────────────────────────────────────────────────

BACKUP_ARCHIVE = "backup.tar.gz"
RESTORE_ENDPOINT = "/api/v1/backup/restore"
BACKUP_ENDPOINT = "/api/v1/backup"
HEALTH_ENDPOINT = "/health"
HF_API_SPACES = "https://huggingface.co/api/spaces"


def _build_url(base: str, path: str) -> str:
    base = base.rstrip("/")
    path = path.lstrip("/")
    return f"{base}/{path}"


def http_get(url: str, headers: dict[str, str], timeout: int = 120) -> tuple[int, bytes]:
    req = urllib.request.Request(url, headers=headers, method="GET")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status, resp.read()


def http_get_stream(
    url: str, headers: dict[str, str], dest: Path, timeout: int = 300
) -> tuple[int, int]:
    req = urllib.request.Request(url, headers=headers, method="GET")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        total = 0
        with open(dest, "wb") as f:
            while True:
                chunk = resp.read(65536)
                if not chunk:
                    break
                f.write(chunk)
                total += len(chunk)
        return resp.status, total


def http_post_multipart(
    url: str,
    file_path: Path,
    field_name: str,
    headers: dict[str, str],
    timeout: int = 300,
) -> tuple[int, bytes]:
    boundary = "----DeployBoundary" + hex(int(time.time() * 1e6))[2:]
    data = io.BytesIO()

    data.write(f"--{boundary}\r\n".encode())
    data.write(
        f'Content-Disposition: form-data; name="{field_name}"; filename="{file_path.name}"\r\n'.encode()
    )
    data.write(b"Content-Type: application/gzip\r\n\r\n")
    data.write(file_path.read_bytes())
    data.write(f"\r\n--{boundary}--\r\n".encode())

    body = data.getvalue()
    content_type = f"multipart/form-data; boundary={boundary}"

    req_headers = {**headers, "Content-Type": content_type}
    req = urllib.request.Request(url, data=body, headers=req_headers, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status, resp.read()


def http_get_json(
    url: str, headers: dict[str, str], timeout: int = 30
) -> Optional[dict[str, Any]]:
    try:
        req = urllib.request.Request(url, headers=headers, method="GET")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except Exception:
        return None


def http_post_json(
    url: str, headers: dict[str, str], data: Optional[bytes] = None, timeout: int = 30
) -> Optional[int]:
    try:
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status
    except urllib.error.HTTPError as e:
        return e.code
    except Exception:
        return None


def http_head(url: str, timeout: int = 30) -> Optional[int]:
    try:
        req = urllib.request.Request(url, method="HEAD")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status
    except urllib.error.URLError:
        return None


# ─────────────────────────────────────────────────────────────
# Git helpers
# ─────────────────────────────────────────────────────────────


def _run_git(args: list[str], cwd: str | None = None) -> tuple[int, str]:
    cmd = ["git"] + args
    result = subprocess.run(cmd, capture_output=True, text=True, cwd=cwd)
    return result.returncode, result.stdout.strip()


def git_status(cwd: str) -> list[str]:
    rc, out = _run_git(["status", "--porcelain"], cwd=cwd)
    if rc != 0:
        return []
    lines = [line for line in out.split("\n") if line.strip()]
    return lines


def git_add_all(cwd: str) -> bool:
    rc, _ = _run_git(["add", "-A", "."], cwd=cwd)
    return rc == 0


def git_commit(cwd: str, message: str) -> tuple[bool, str]:
    rc, out = _run_git(["commit", "-m", message], cwd=cwd)
    if rc == 0:
        m = re.search(r"\[[^\]]+ ([a-f0-9]+)\]", out)
        sha = m.group(1) if m else "unknown"
        return True, sha
    if "nothing to commit" in out.lower() or "no changes" in out.lower():
        return True, "no-change"
    return False, out


def git_push(cwd: str, remote: str = "origin", branch: str = "main") -> tuple[bool, str]:
    rc, out = _run_git(["push", remote, branch], cwd=cwd)
    return rc == 0, out


def _parse_hf_space_from_url(url: str) -> tuple[str, str]:
    m = re.search(r"huggingface\.co/spaces/([^/]+)/([^/\s]+)", url)
    if m:
        return m.group(1), m.group(2).rstrip("/")
    return "aetherbase", "llm-ready-data"


def git_remote_url(cwd: str) -> str:
    rc, out = _run_git(["remote", "get-url", "origin"], cwd=cwd)
    if rc != 0:
        return ""
    return out


def git_log(cwd: str, n: int = 3) -> str:
    _, out = _run_git(["log", "--oneline", f"-{n}"], cwd=cwd)
    return out


# ─────────────────────────────────────────────────────────────
# Core deployment logic
# ─────────────────────────────────────────────────────────────


class AuditReport:
    def __init__(self):
        self.start_time = time.time()
        self.fields: dict[str, Any] = {
            "backup_started": "",
            "backup_completed": "",
            "backup_files": 0,
            "backup_size": 0,
            "backup_status": "pending",
            "commit_status": "pending",
            "commit_hash": "",
            "push_status": "pending",
            "deploy_status": "pending",
            "deploy_duration": "",
            "restore_status": "pending",
            "restore_files": 0,
            "verification_status": "pending",
            "errors": [],
            "warnings": [],
        }

    def end(self):
        self.fields["total_duration"] = _human_duration(time.time() - self.start_time)

    def set(self, key: str, value: Any):
        self.fields[key] = value

    def error(self, msg: str):
        self.fields["errors"].append(msg)

    def warn(self, msg: str):
        self.fields["warnings"].append(msg)

    def _row(self, label: str, status: str, detail: str) -> str:
        icon = _indicator(status)
        return f"{_TABLE_MV} {label:<28} {icon:<2}    {detail:<30} {_TABLE_MV}"

    def _sep(self) -> str:
        return f"{_TABLE_ML}{_TABLE_MH * 28}{_TABLE_MM}{_TABLE_MH * 4}{_TABLE_MM}{_TABLE_MH * 32}{_TABLE_MR}"

    def print(self):
        self.end()
        f = self.fields
        err_count = len(f["errors"])
        warn_count = len(f["warnings"])

        start_dt = datetime.fromtimestamp(self.start_time).strftime("%c")

        lines: list[str] = []
        lines.append("")
        lines.append(
            f"{_TABLE_TL}{_TABLE_H * 16}{_TABLE_H * 16}{_TABLE_H * 18}{_TABLE_H * 20}"
            f"{_TABLE_TR}"
        )

        title = "DEPLOYMENT AUDIT REPORT"
        padding = (70 - len(title)) // 2
        lines.append(
            f"{_TABLE_V} {' ' * padding}{_bold(title)}{' ' * (70 - padding - len(title))} {_TABLE_V}"
        )
        lines.append(
            f"{_TABLE_V} {' ' * (70 - len(start_dt))}{start_dt} {_TABLE_V}"
        )
        lines.append(
            f"{_TABLE_TM}{_TABLE_H * 70}{_TABLE_H * 0}"
        )

        lines.append(self._row("Category", "", "Status | Detail"))
        lines.append(
            f"{_TABLE_TML}{_TABLE_MH * 70}"
        )

        lines.append(self._row("", "", ""))

        # ── Backup section ──
        lines.append(
            self._row(
                "Backup Started",
                "passed" if f["backup_started"] else "failed",
                f["backup_started"] or "\u2014",
            )
        )
        lines.append(
            self._row(
                "Backup Completed",
                "passed" if f["backup_completed"] else "failed",
                f["backup_completed"] or "\u2014",
            )
        )
        lines.append(
            self._row(
                "Files Backed Up",
                "passed" if f["backup_files"] > 0 else "failed",
                str(f["backup_files"]),
            )
        )
        lines.append(
            self._row(
                "Total Backup Size",
                "passed" if f["backup_size"] > 0 else "failed",
                _human_size(f["backup_size"]),
            )
        )

        lines.append(self._sep())

        # ── Deployment section ──
        lines.append(
            self._row(
                "Commit Status",
                f["commit_status"],
                f["commit_hash"] or "\u2014",
            )
        )
        lines.append(
            self._row(
                "Push Status",
                f["push_status"],
                _ok("Pushed") if f["push_status"] == "passed" else _fail("Failed"),
            )
        )
        dep_icon = "passed" if f["deploy_status"] == "passed" else "failed"
        lines.append(
            self._row(
                "HF Space Deployment",
                dep_icon,
                f["deploy_duration"] or "\u2014",
            )
        )

        lines.append(self._sep())

        # ── Restore section ──
        lines.append(
            self._row(
                "Data Restoration",
                f["restore_status"],
                _ok("Restored") if f["restore_status"] == "passed" else _fail("Failed"),
            )
        )
        lines.append(
            self._row(
                "Files Restored",
                "passed" if f["restore_files"] > 0 else "failed",
                str(f["restore_files"]),
            )
        )
        lines.append(
            self._row(
                "Verification Status",
                f["verification_status"],
                _ok("Verified") if f["verification_status"] == "passed" else _fail("Failed"),
            )
        )

        lines.append(self._sep())

        # ── Summary section ──
        lines.append(
            self._row(
                "Total Execution Time",
                "passed",
                f.get("total_duration", ""),
            )
        )
        lines.append(
            self._row(
                "Errors",
                "failed" if err_count > 0 else "passed",
                str(err_count),
            )
        )
        lines.append(
            self._row(
                "Warnings",
                "passed" if warn_count == 0 else "skipped",
                str(warn_count),
            )
        )

        lines.append(
            f"{_TABLE_BL}{_TABLE_H * 16}{_TABLE_H * 16}{_TABLE_H * 18}{_TABLE_H * 20}"
            f"{_TABLE_BR}"
        )

        for line in lines:
            print(line)

        if err_count > 0:
            print(f"\n{_fail('Errors:')}")
            for e in f["errors"]:
                print(f"  {_fail(_BULLET)} {e}")

        if warn_count > 0:
            print(f"\n{_warn('Warnings:')}")
            for w in f["warnings"]:
                print(f"  {_warn(_BULLET)} {w}")


def run_deployment(args: argparse.Namespace, env: dict[str, str]) -> int:
    audit = AuditReport()
    cwd = os.getcwd()
    persistence_dir = Path(cwd) / "persistence"
    persistence_dir.mkdir(parents=True, exist_ok=True)

    space_url = args.space_url or env.get("SPACE_URL", "https://aetherbase-llm-ready-data.hf.space")
    api_key = args.api_key or env.get("api_key", env.get("API_KEY", ""))
    hf_token = args.hf_token or env.get("hf_token", env.get("HF_TOKEN", ""))
    skip_backup = args.skip_backup

    if not api_key:
        print(f"  {_fail('ERROR:')} No API key found. Set API_KEY in .env or pass --api-key.")
        return 1

    if not hf_token:
        print(f"  {_warn('WARNING:')} No HF token. Can't verify Space build status. Set HF_TOKEN in .env or pass --hf-token.")
        hf_token = None

    auth_headers = {
        "Authorization": f"Bearer {api_key}",
        "User-Agent": "deploy.py/1.0",
    }

    remote_url = git_remote_url(cwd)
    owner, space_name = _parse_hf_space_from_url(remote_url) if remote_url else ("aetherbase", "llm-ready-data")
    hf_api_headers = None
    if hf_token:
        hf_api_headers = {
            "Authorization": f"Bearer {hf_token}",
            "User-Agent": "deploy.py/1.0",
        }

    # ──────────────────────────────────────────────
    # PHASE 1: LOCK WRITES + BACKUP
    # ──────────────────────────────────────────────
    print(f"\n{_bold('PHASE 1/4: Lock Writes & Backup')}")
    print(f"  {_INFO_ICON} Backing up data from: {space_url}")

    if skip_backup:
        print(f"  {_WARN_ICON} Backup skipped (--skip-backup)")
        audit.set("backup_status", "skipped")
    else:
        # Enable maintenance mode to freeze data state
        maint_url = _build_url(space_url, "/api/v1/system/maintenance/enable")
        maint_enabled = False
        try:
            status = http_post_json(maint_url, auth_headers, timeout=15)
            if status and status == 200:
                maint_enabled = True
                print(f"  {_OK_ICON} Maintenance enabled — writes blocked")
        except Exception:
            print(f"  {_WARN_ICON} Could not enable maintenance (endpoint may not exist yet)")
            print(f"  {_WARN_ICON} Data may be lost if writes occur during deployment")

        backup_file = persistence_dir / BACKUP_ARCHIVE
        backup_url = _build_url(space_url, BACKUP_ENDPOINT)
        audit.set("backup_started", _timestamp())
        backed_up = False

        # Try remote backup endpoint first
        try:
            status, total_bytes = http_get_stream(backup_url, auth_headers, backup_file, timeout=120)
            if status == 200:
                backed_up = True
                print(f"  {_OK_ICON} Remote backup successful")
        except Exception as exc:
            print(f"  {_WARN_ICON} Remote backup unavailable: {exc}")

        # Fall back to local backup
        if not backed_up:
            data_dir = Path(cwd) / "data"
            if data_dir.is_dir():
                print(f"  {_INFO_ICON} Falling back to local backup: {data_dir}")
                try:
                    buf = io.BytesIO()
                    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
                        for path in sorted(data_dir.rglob("*")):
                            if path.is_file():
                                arcname = path.relative_to(data_dir.parent)
                                tar.add(str(path), arcname=str(arcname))
                    buf.seek(0)
                    backup_file.write_bytes(buf.read())

                    total_bytes = backup_file.stat().st_size
                    file_count = 0
                    with tarfile.open(backup_file, "r:gz") as tar:
                        file_count = sum(1 for m in tar.getmembers() if m.isfile())

                    backed_up = True
                    print(f"  {_OK_ICON} Local backup saved: {backup_file}")
                except Exception as exc:
                    msg = f"Local backup failed: {exc}"
                    audit.error(msg)
                    print(f"  {_FAIL_ICON} {msg}")
                    audit.print()
                    return 1
            else:
                msg = f"No data found locally ({data_dir}) and remote backup unavailable"
                audit.warn(msg)
                print(f"  {_WARN_ICON} {msg}")
                audit.set("backup_files", 0)
                audit.set("backup_size", 0)
                audit.set("backup_completed", _timestamp())
                audit.set("backup_status", "skipped")

        if backed_up:
            file_count = 0
            try:
                with tarfile.open(backup_file, "r:gz") as tar:
                    file_count = sum(1 for m in tar.getmembers() if m.isfile())
            except Exception:
                audit.warn("Could not count files in backup archive")
                file_count = backup_file.stat().st_size

            audit.set("backup_completed", _timestamp())
            audit.set("backup_files", file_count)
            audit.set("backup_size", total_bytes)
            audit.set("backup_status", "passed")
            print(f"  {_OK_ICON} Files: {file_count} | Size: {_human_size(total_bytes)}")
            if maint_enabled:
                print(f"  {_OK_ICON} Writes frozen — backup is consistent")

    # ──────────────────────────────────────────────
    # PHASE 2: GIT COMMIT & PUSH
    # ──────────────────────────────────────────────
    print(f"\n{_bold('PHASE 2/4: Git Commit & Push')}")

    changed = git_status(cwd)
    if not changed:
        print(f"  {_WARN_ICON} No changes to commit (working tree clean)")
        audit.set("commit_status", "passed")
        audit.set("commit_hash", "no-change")
        audit.set("push_status", "passed")
    else:
        print(f"  {_INFO_ICON} Staging {len(changed)} file(s)...")
        if not git_add_all(cwd):
            msg = "Git add failed"
            audit.error(msg)
            audit.set("commit_status", "failed")
            print(f"  {_FAIL_ICON} {msg}")
            audit.print()
            return 1

        commit_msg = args.message or f"deploy: auto-deploy {_timestamp()}"
        print(f"  {_INFO_ICON} Committing: {commit_msg}")
        ok, sha = git_commit(cwd, commit_msg)
        if not ok:
            msg = f"Git commit failed: {sha}"
            audit.error(msg)
            audit.set("commit_status", "failed")
            print(f"  {_FAIL_ICON} {msg}")
            audit.print()
            return 1

        audit.set("commit_status", "passed")
        audit.set("commit_hash", sha)
        print(f"  {_OK_ICON} Committed: {sha}")

        print(f"  {_INFO_ICON} Pushing to origin/main...")
        ok, push_out = git_push(cwd)
        if not ok:
            msg = f"Git push failed"
            audit.error(msg)
            audit.set("push_status", "failed")
            print(f"  {_FAIL_ICON} {msg}")
            print(f"  {push_out}")
            audit.print()
            return 1

        audit.set("push_status", "passed")
        print(f"  {_OK_ICON} Push successful")

    # ──────────────────────────────────────────────
    # PHASE 3: WAIT FOR REDEPLOYMENT (via HF API)
    # ──────────────────────────────────────────────
    print(f"\n{_bold('PHASE 3/4: Waiting for HF Space Redeployment')}")

    api_url = f"{HF_API_SPACES}/{owner}/{space_name}"
    deploy_start = time.time()
    max_wait = args.timeout
    poll_interval = 10
    waited = 0
    deployed = False
    last_error = ""
    seen_building = False

    print(f"  {_INFO_ICON} Tracking {owner}/{space_name} via HF API...")

    if hf_api_headers:
        initial = http_get_json(api_url, hf_api_headers, timeout=15) or {}
        prev_stage = initial.get("runtime", {}).get("stage", "")

        while waited < max_wait:
            data = http_get_json(api_url, hf_api_headers, timeout=15)
            stage = (data or {}).get("runtime", {}).get("stage", "")

            if stage == "BUILDING":
                if not seen_building:
                    print(f"  {_INFO_ICON} Build started (stage: BUILDING)")
                    seen_building = True

            elif stage == "RUNNING":
                if seen_building:
                    print(f"  {_OK_ICON} Build complete (stage: RUNNING)")
                    deployed = True
                    break
                if waited > 30 and prev_stage == "RUNNING":
                    print(f"  {_OK_ICON} Space is running (may have skipped BUILDING stage)")
                    deployed = True
                    break

            elif stage in ("PAUSED", "STOPPED", "NO_APP"):
                last_error = f"Space in unexpected state: {stage}"
                deployed = False
                break

            prev_stage = stage
            time.sleep(poll_interval)
            waited += poll_interval
            if waited % 30 == 0 and not seen_building:
                print(f"  {_INFO_ICON} Still waiting for build to start... ({waited}s)")
            elif waited % 30 == 0 and seen_building:
                print(f"  {_INFO_ICON} Still building... ({waited}s)")

    deploy_duration = time.time() - deploy_start

    if not hf_api_headers:
        print(f"  {_WARN_ICON} No HF token — assuming deployment is underway")
        print(f"  {_INFO_ICON} Waiting {max_wait}s for build + startup...")
        time.sleep(max_wait)
        deployed = True

    if deployed:
        health_url = _build_url(space_url, HEALTH_ENDPOINT)
        try:
            req = urllib.request.Request(health_url, method="GET")
            with urllib.request.urlopen(req, timeout=15) as resp:
                if resp.status == 200:
                    print(f"  {_OK_ICON} App health check passed (new instance serving)")
        except Exception:
            print(f"  {_WARN_ICON} App health check unavailable (may still be starting)")

        audit.set("deploy_status", "passed")
        audit.set("deploy_duration", _human_duration(deploy_duration))
        print(f"  {_OK_ICON} Space redeployed in {_human_duration(deploy_duration)}")
    else:
        audit.set("deploy_status", "failed")
        audit.set("deploy_duration", _human_duration(deploy_duration))
        msg = f"Space did not redeploy within {max_wait}s (last stage: {stage}, error: {last_error})"
        audit.error(msg)
        print(f"  {_FAIL_ICON} {msg}")
        audit.print()
        return 1

    # ──────────────────────────────────────────────
    # PHASE 4: DATA RESTORE
    # ──────────────────────────────────────────────
    print(f"\n{_bold('PHASE 4/4: Data Restoration')}")

    if skip_backup:
        print(f"  {_WARN_ICON} Restore skipped (no backup)")
        audit.set("restore_status", "skipped")
        audit.set("verification_status", "skipped")
        audit.set("restore_files", 0)
    else:
        backup_file = persistence_dir / BACKUP_ARCHIVE
        if not backup_file.is_file():
            msg = f"Backup file not found: {backup_file}"
            audit.error(msg)
            audit.set("restore_status", "failed")
            print(f"  {_FAIL_ICON} {msg}")
            audit.print()
            return 1

        restore_url = _build_url(space_url, RESTORE_ENDPOINT)
        print(f"  {_INFO_ICON} Restoring data to: {space_url}")

        try:
            status, body = http_post_multipart(
                restore_url, backup_file, "file", auth_headers, timeout=300
            )

            if status == 200:
                try:
                    result = json.loads(body)
                    restored = result.get("files_restored", 0)
                except Exception:
                    restored = 0

                audit.set("restore_files", max(restored, 1))
                audit.set("restore_status", "passed")
                print(f"  {_OK_ICON} Data restored: {restored} files")

                # Verification
                print(f"  {_INFO_ICON} Verifying data restoration...")
                try:
                    verify_backup = persistence_dir / "verify_check.tar.gz"
                    status_v, _ = http_get_stream(
                        _build_url(space_url, BACKUP_ENDPOINT),
                        auth_headers,
                        verify_backup,
                        timeout=120,
                    )

                    if status_v == 200:
                        with tarfile.open(verify_backup, "r:gz") as tar:
                            verify_files = sum(1 for m in tar.getmembers() if m.isfile())
                        verify_backup.unlink(missing_ok=True)

                        if verify_files >= audit.fields.get("backup_files", 0) * 0.9:
                            audit.set("verification_status", "passed")
                            print(f"  {_OK_ICON} Verification passed: {verify_files} files found")
                        else:
                            audit.set("verification_status", "failed")
                            audit.warn(
                                f"Verification file count ({verify_files}) mismatch with backup ({audit.fields.get('backup_files', 0)})"
                            )
                            print(f"  {_WARN_ICON} Verification: file count mismatch ({verify_files} vs {audit.fields.get('backup_files', 0)})")
                    else:
                        audit.set("verification_status", "failed")
                        audit.warn(f"Verification request returned HTTP {status_v}")
                        print(f"  {_WARN_ICON} Verification failed (HTTP {status_v})")
                except Exception as exc:
                    audit.set("verification_status", "failed")
                    audit.warn(f"Verification error: {exc}")
                    print(f"  {_WARN_ICON} Verification error: {exc}")
            else:
                audit.set("restore_status", "failed")
                msg = f"Restore returned HTTP {status}: {body[:200]}"
                audit.error(msg)
                print(f"  {_FAIL_ICON} {msg}")
                audit.print()
                return 1

        except Exception as exc:
            audit.set("restore_status", "failed")
            msg = f"Restore failed: {exc}"
            audit.error(msg)
            print(f"  {_FAIL_ICON} {msg}")
            audit.print()
            return 1

    # ──────────────────────────────────────────────
    # SUMMARY
    # ──────────────────────────────────────────────
    print(f"\n{_bold('=' * 70)}")
    audit.print()

    has_errors = len(audit.fields["errors"]) > 0
    if has_errors:
        print(f"\n{_fail('Deployment completed with errors.')}")
        return 1
    else:
        print(f"\n{_ok('Deployment completed successfully.')}")
        return 0


# ─────────────────────────────────────────────────────────────
# CLI entrypoint
# ─────────────────────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(
        description="Deploy to Hugging Face Spaces with persistent data handling.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Environment variables (from .env or process env):\n"
            "  API_KEY      API key for the Hugging Face Space\n"
            "  SPACE_URL    Base URL of the deployed Space\n"
            "  HF_TOKEN     Hugging Face API token (for build tracking)\n"
        ),
    )
    parser.add_argument(
        "--space-url",
        default="",
        help="Base URL of the Hugging Face Space (e.g. https://aetherbase-llm-ready-data.hf.space)",
    )
    parser.add_argument(
        "--api-key",
        default="",
        help="API key for authenticating with the Space",
    )
    parser.add_argument(
        "--message", "-m",
        default="",
        help="Git commit message",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=300,
        help="Maximum wait time (seconds) for Space redeployment (default: 300)",
    )
    parser.add_argument(
        "--hf-token",
        default="",
        help="Hugging Face API token (for tracking Space build status)",
    )
    parser.add_argument(
        "--skip-backup",
        action="store_true",
        help="Skip the backup and restore phases",
    )
    parser.add_argument(
        "--env-file",
        default=".env",
        help="Path to .env file (default: .env)",
    )

    args = parser.parse_args()
    env = load_env(args.env_file)

    # Merge with process environment (process env takes precedence)
    for key in ("API_KEY", "SPACE_URL", "HF_TOKEN"):
        if os.environ.get(key):
            env[key] = os.environ[key]

    try:
        rc = run_deployment(args, env)
    except KeyboardInterrupt:
        print(f"\n{_warn('Deployment interrupted by user.')}")
        rc = 130
    except Exception as exc:
        print(f"\n{_fail('Unexpected error: ' + str(exc))}")
        rc = 1

    sys.exit(rc)


if __name__ == "__main__":
    main()