#!/usr/bin/env python3
"""Roko Backup Script — refactored v2.

Collect -> archive -> upload (streaming progress) -> verify -> clean -> report.
Runs from system crontab, standalone Python 3.10+.
"""

from __future__ import annotations

import abc
import argparse
import collections.abc
import contextlib
import datetime
import json
import logging
import logging.handlers
import os
import re
import signal
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

log = logging.getLogger("roko-backup")


# --- Exceptions ---------------------------------------------------------------


class BackupError(Exception):
    """Base exception for backup pipeline failures."""


class LockError(BackupError):
    """Raised when the lock cannot be acquired."""


# --- Shared rclone environment ------------------------------------------------


def _rclone_env() -> dict[str, str]:
    """Return environment dict with NO_PROXY set for Yandex Cloud."""
    return {
        **os.environ,
        "NO_PROXY": "storage.yandexcloud.net,yandexcloud.net",
    }


# --- Config -------------------------------------------------------------------


@dataclass(frozen=True)
class Config:
    """All paths and limits for the backup pipeline."""

    home: Path = field(default_factory=lambda: Path.home())
    report_file: Path = field(
        default_factory=lambda: Path.home() / ".openclaw" / "workspace" / "roko-backup-report.json"
    )
    archive_dir: Path = field(default_factory=lambda: Path.home() / "backups")
    log_file: Path = field(default_factory=lambda: Path.home() / "backups" / "roko-backup.log")
    rclone_remote: str = "yabackup:for-backup/roko"
    max_cloud_backups: int = 14
    max_local_backups: int = 3
    max_log_size: int = 1_048_576  # 1 MiB
    max_log_backups: int = 5
    lock_path: Path = field(default_factory=lambda: Path("/tmp/roko-backup.lock"))

    @property
    def workspace(self) -> Path:
        """Return the workspace directory."""
        return self.home / ".openclaw" / "workspace"

    @property
    def config_file(self) -> Path:
        """Return the main openclaw config file."""
        return self.home / ".openclaw" / "openclaw.json"

    @property
    def scripts_dir(self) -> Path:
        """Return the scripts directory."""
        return self.home / ".openclaw" / "scripts"

    @property
    def media_dir(self) -> Path:
        """Return the media directory."""
        return self.home / ".openclaw" / "media"

    @property
    def plugin_skills_dir(self) -> Path:
        """Return the plugin-skills directory."""
        return self.home / ".openclaw" / "plugin-skills"

    @property
    def agent_db(self) -> Path:
        """Return the agent SQLite database path."""
        return self.home / ".openclaw" / "agents" / "main" / "agent" / "openclaw-agent.sqlite"

    @classmethod
    def from_env(cls) -> Config:
        """Build config from environment variables with built-in defaults."""
        home = Path.home()
        return cls(
            home=home,
            report_file=Path(os.environ.get("ROKO_REPORT_FILE", str(home / ".openclaw/workspace/roko-backup-report.json"))),
            archive_dir=Path(os.environ.get("ROKO_ARCHIVE_DIR", str(home / "backups"))),
            log_file=Path(os.environ.get("ROKO_LOG_FILE", str(home / "backups/roko-backup.log"))),
            rclone_remote=os.environ.get("ROKO_RCLONE_REMOTE", "yabackup:for-backup/roko"),
            max_cloud_backups=int(os.environ.get("ROKO_MAX_CLOUD", "14")),
            max_local_backups=int(os.environ.get("ROKO_MAX_LOCAL", "3")),
        )

    @classmethod
    def from_file(cls, path: Path) -> Config:
        """Load config from a JSON file, falling back to defaults."""
        data: dict[str, Any] = {}
        if path.exists():
            data = json.loads(path.read_text())
        home = Path(data.get("home", str(Path.home())))
        return cls(
            home=home,
            report_file=Path(data.get("report_file", str(home / ".openclaw/workspace/roko-backup-report.json"))),
            archive_dir=Path(data.get("archive_dir", str(home / "backups"))),
            log_file=Path(data.get("log_file", str(home / "backups/roko-backup.log"))),
            rclone_remote=data.get("rclone_remote", "yabackup:for-backup/roko"),
            max_cloud_backups=data.get("max_cloud_backups", 14),
            max_local_backups=data.get("max_local_backups", 3),
            max_log_size=data.get("max_log_size", 1_048_576),
            max_log_backups=data.get("max_log_backups", 5),
        )


# --- LockFile -----------------------------------------------------------------


class LockFile:
    """Context manager that prevents concurrent backup runs via a PID file."""

    def __init__(self, path: Path) -> None:
        self.path = path
        self._acquired = False

    def __enter__(self) -> LockFile:
        if not self.acquire():
            raise LockError(f"Another backup is already running (lock: {self.path})")
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,
    ) -> None:
        self.release()

    def acquire(self) -> bool:
        """Acquire the lock. Returns False if another live process holds it."""
        if self.path.exists():
            try:
                pid = int(self.path.read_text().strip())
                os.kill(pid, 0)
                return False
            except (ValueError, ProcessLookupError, PermissionError, OSError):
                log.warning("Stale lock file found (PID dead), removing")
                self.path.unlink(missing_ok=True)

        self.path.write_text(str(os.getpid()))
        self._acquired = True
        signal.signal(signal.SIGTERM, self._signal_handler)
        signal.signal(signal.SIGINT, self._signal_handler)
        return True

    def release(self) -> None:
        """Release the lock and remove the file."""
        if self._acquired:
            self.path.unlink(missing_ok=True)
            self._acquired = False

    def _signal_handler(self, signum: int, _frame: Any) -> None:
        self.release()
        sys.exit(128 + signum)


# --- Collector ----------------------------------------------------------------


class Collector:
    """Gathers files to back up into a staging directory."""

    EXCLUDE_PATTERNS: tuple[str, ...] = (
        "node_modules",
        "__pycache__",
        "*.pyc",
        ".git",
        "roko-pulse.jsonl",
        "cache",
        "*.tmp",
        ".cache",
    )

    def __init__(self, config: Config) -> None:
        self.config = config

    def collect(self, staging: Path) -> tuple[list[tuple[Path, Path]], int]:
        """Copy workspace/config/scripts/media/plugin-skills/agent-db into staging.

        Returns (list of (source, dest) pairs, total bytes collected).
        """
        collected: list[tuple[Path, Path]] = []
        total_bytes = 0
        ignore = shutil.ignore_patterns(*self.EXCLUDE_PATTERNS)

        targets: list[tuple[Path, str]] = [
            (self.config.workspace, "openclaw/workspace"),
            (self.config.config_file, "openclaw/openclaw.json"),
            (self.config.scripts_dir, "openclaw/scripts"),
            (self.config.media_dir, "openclaw/media"),
            (self.config.plugin_skills_dir, "openclaw/plugin-skills"),
            (self.config.agent_db, "openclaw/openclaw-agent.sqlite"),
        ]

        for src, dest_rel in targets:
            dest = staging / dest_rel
            if not src.exists():
                log.warning("Source not found, skipping: %s", src)
                continue

            dest.parent.mkdir(parents=True, exist_ok=True)

            try:
                if src.is_dir():
                    shutil.copytree(src, dest, ignore=ignore, dirs_exist_ok=True)
                else:
                    shutil.copy2(src, dest)
            except Exception as exc:
                log.warning("Failed to copy %s: %s", src, exc)
                continue

            size = self._tree_size(dest)
            collected.append((src, dest))
            total_bytes += size
            log.info("Collected %s -> %s (%d bytes)", src.name, dest_rel, size)

        return collected, total_bytes

    @staticmethod
    def _tree_size(path: Path) -> int:
        """Return total size in bytes for a file or directory."""
        if path.is_file():
            return path.stat().st_size
        return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())


# --- Archiver -----------------------------------------------------------------


class Archiver:
    """Creates a tar.gz archive from a staging directory."""

    def __init__(self, config: Config) -> None:
        self.config = config

    def archive(self, staging: Path, backup_name: str) -> tuple[Path, int]:
        """Create tar.gz archive. Returns (archive_path, size_bytes)."""
        self.config.archive_dir.mkdir(parents=True, exist_ok=True)
        archive_path = self.config.archive_dir / f"{backup_name}.tar.gz"

        log.info("Creating archive: %s", archive_path)
        with tarfile.open(archive_path, mode="w:gz") as tar:
            tar.add(staging, arcname=".")

        size = archive_path.stat().st_size
        log.info("Archive created: %s (%d bytes)", archive_path.name, size)
        return archive_path, size


# --- Uploader -----------------------------------------------------------------


class Uploader:
    """Uploads archive via rclone with streaming progress and retry."""

    BACKOFF_SCHEDULE: tuple[int, ...] = (10, 30, 60)
    UPLOAD_PROGRESS_INTERVAL: int = 30  # seconds between progress logs
    PROGRESS_PERCENT_STEP: int = 5  # log every N percent

    _RE_TRANSFERRED = re.compile(
        r"Transferred:\s+[\w.]+\s+([\d.]+\s*\w+)\s+/\s+([\d.]+\s*\w+),\s+(\d+)%"
    )

    def __init__(self, config: Config) -> None:
        self.config = config

    def upload(self, archive_path: Path) -> bool:
        """Upload archive to remote with streaming progress. Returns True on success."""
        max_attempts = len(self.BACKOFF_SCHEDULE) + 1

        for attempt in range(1, max_attempts + 1):
            log.info("Upload attempt %d/%d", attempt, max_attempts)

            if self._upload_once(archive_path):
                return True

            if attempt < max_attempts:
                delay = self.BACKOFF_SCHEDULE[min(attempt - 1, len(self.BACKOFF_SCHEDULE) - 1)]
                log.info("Retrying in %ds...", delay)
                time.sleep(delay)

        log.error("Upload failed after %d attempts", max_attempts)
        return False

    def _upload_once(self, archive_path: Path) -> bool:
        """Run a single rclone copy with streaming stderr.

        Parses progress lines from rclone and logs periodically.
        """
        cmd = [
            "rclone", "copy",
            str(archive_path),
            f"{self.config.rclone_remote}/",
            "--progress",
        ]
        last_log_ts: float = 0.0
        last_pct: int = -1

        with subprocess.Popen(
            cmd,
            stderr=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            text=True,
            env=_rclone_env(),
        ) as proc:
            assert proc.stderr is not None
            for line in proc.stderr:
                line = line.rstrip("\n")
                if not line:
                    continue

                m = self._RE_TRANSFERRED.search(line)
                if m:
                    transferred, total, pct_str = m.group(1), m.group(2), int(m.group(3))
                    now = time.monotonic()
                    new_pct_bucket = pct_str // self.PROGRESS_PERCENT_STEP * self.PROGRESS_PERCENT_STEP
                    if now - last_log_ts >= self.UPLOAD_PROGRESS_INTERVAL or new_pct_bucket > last_pct:
                        log.info("Upload: %d%% (%s / %s)", pct_str, transferred, total)
                        last_log_ts = now
                        last_pct = new_pct_bucket
                else:
                    log.debug("rclone: %s", line)

            proc.wait()

        if proc.returncode == 0:
            log.info("Upload succeeded")
            return True

        log.warning("Upload failed (rc=%d)", proc.returncode)
        return False


# --- Verifier -----------------------------------------------------------------


class Verifier:
    """Confirms the uploaded file exists in the remote."""

    def __init__(self, config: Config) -> None:
        self.config = config

    def verify(self, backup_name: str) -> tuple[bool, str]:
        """Check that backup_name.tar.gz is listed in the remote.

        Returns (success, details).
        """
        result = subprocess.run(
            ["rclone", "ls", self.config.rclone_remote],
            capture_output=True,
            text=True,
            check=False,
            env=_rclone_env(),
        )

        if result.returncode != 0:
            return False, f"rclone ls failed (rc={result.returncode}): {result.stderr[:200]}"

        target = f"{backup_name}.tar.gz"
        if target in result.stdout:
            return True, f"Verified: {target} found in remote"

        return False, f"Not found in remote: {target}"


# --- Cleaner ------------------------------------------------------------------


class Cleaner:
    """Manages retention of local and cloud backups."""

    def __init__(self, config: Config) -> None:
        self.config = config

    def clean_local(self) -> int:
        """Delete local archives beyond max_local_backups. Returns count deleted."""
        archives = sorted(
            self.config.archive_dir.glob("roko-*.tar.gz"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        deleted = 0
        for old in archives[self.config.max_local_backups :]:
            log.info("Deleting local: %s", old.name)
            old.unlink()
            deleted += 1
        return deleted

    def clean_cloud(self) -> int:
        """Delete cloud archives beyond max_cloud_backups. Returns count deleted."""
        result = subprocess.run(
            ["rclone", "ls", self.config.rclone_remote],
            capture_output=True,
            text=True,
            check=False,
            env=_rclone_env(),
        )
        if result.returncode != 0:
            log.warning("Cannot list cloud backups for cleanup: %s", result.stderr[:200])
            return 0

        entries: list[tuple[int, str]] = []
        for line in result.stdout.strip().splitlines():
            parts = line.split(None, 1)
            if len(parts) == 2:
                try:
                    entries.append((int(parts[0]), parts[1]))
                except ValueError:
                    continue

        entries.sort(key=lambda e: e[1])
        deleted = 0
        to_remove = entries[: max(0, len(entries) - self.config.max_cloud_backups)]

        for _size, name in to_remove:
            log.info("Deleting cloud: %s", name)
            subprocess.run(
                ["rclone", "delete", f"{self.config.rclone_remote}/{name}"],
                capture_output=True,
                text=True,
                check=False,
                env=_rclone_env(),
            )
            deleted += 1

        return deleted


# --- Reporter -----------------------------------------------------------------


class Reporter:
    """Writes JSON backup report to the workspace."""

    def __init__(self, config: Config) -> None:
        self.config = config

    def write(
        self,
        status: str,
        name: str | None = None,
        size_bytes: int | None = None,
        error: str | None = None,
    ) -> None:
        """Write the report JSON file."""
        report = {
            "timestamp": datetime.datetime.now().astimezone().isoformat(),
            "status": status,
            "name": name,
            "size_bytes": size_bytes,
            "error": error,
        }
        self.config.report_file.parent.mkdir(parents=True, exist_ok=True)
        self.config.report_file.write_text(json.dumps(report, indent=2) + "\n")
        log.info("Report written: %s (%s)", self.config.report_file, status)


# --- Pipeline -----------------------------------------------------------------


class BackupPipeline:
    """Orchestrates the full backup: collect -> archive -> upload -> verify -> clean -> report."""

    def __init__(self, config: Config, *, dry_run: bool = False) -> None:
        self.config = config
        self.dry_run = dry_run
        self.collector = Collector(config)
        self.archiver = Archiver(config)
        self.uploader = Uploader(config)
        self.verifier = Verifier(config)
        self.cleaner = Cleaner(config)
        self.reporter = Reporter(config)

    def run(self) -> int:
        """Execute the pipeline. Returns exit code."""
        backup_name = f"roko-{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}"
        log.info("Starting backup: %s", backup_name)

        free = shutil.disk_usage(self.config.archive_dir).free
        log.info("Free disk: %s", self._fmt_bytes(free))

        try:
            with tempfile.TemporaryDirectory(prefix="roko-backup-") as tmp:
                staging = Path(tmp)

                # 1. Collect
                t0 = time.monotonic()
                collected, total_bytes = self.collector.collect(staging)
                log.info(
                    "Collection done in %.1fs: %d items, %s",
                    time.monotonic() - t0,
                    len(collected),
                    self._fmt_bytes(total_bytes),
                )

                if self.dry_run:
                    log.info(
                        "[DRY RUN] Would archive and upload %d items (%s)",
                        len(collected),
                        self._fmt_bytes(total_bytes),
                    )
                    for src, dest in collected:
                        log.info("  %s -> %s", src, dest.relative_to(staging))
                    self.reporter.write("dry-run", name=backup_name, size_bytes=total_bytes)
                    return 0

                # 2. Archive
                t0 = time.monotonic()
                archive_path, archive_size = self.archiver.archive(staging, backup_name)
                log.info(
                    "Archiving done in %.1fs: %s",
                    time.monotonic() - t0,
                    self._fmt_bytes(archive_size),
                )

                # 3. Upload
                t0 = time.monotonic()
                if not self.uploader.upload(archive_path):
                    msg = "Upload failed after retries"
                    log.error(msg)
                    self.reporter.write("error", name=backup_name, size_bytes=archive_size, error=msg)
                    return 3
                log.info("Upload done in %.1fs", time.monotonic() - t0)

                # 4. Verify
                ok, detail = self.verifier.verify(backup_name)
                if not ok:
                    log.error("Verification failed: %s", detail)
                    self.reporter.write("error", name=backup_name, size_bytes=archive_size, error=detail)
                    return 3
                log.info("Verification: %s", detail)
        except Exception:
            raise

        # 5. Clean (outside tempdir so archive still exists for local cleanup)
        local_deleted = self.cleaner.clean_local()
        cloud_deleted = self.cleaner.clean_cloud()
        log.info("Cleanup: %d local, %d cloud removed", local_deleted, cloud_deleted)

        # 6. Report
        self.reporter.write("success", name=backup_name, size_bytes=archive_size)
        log.info("Backup complete: %s (%s)", backup_name, self._fmt_bytes(archive_size))
        return 0

    @staticmethod
    def _fmt_bytes(n: int) -> str:
        """Return human-readable byte size."""
        for unit in ("B", "KB", "MB", "GB"):
            if abs(n) < 1024:
                return f"{n:.1f} {unit}"
            n /= 1024  # type: ignore[assignment]
        return f"{n:.1f} TB"


# --- Logging setup ------------------------------------------------------------


def setup_logging(
    *,
    verbose: bool,
    log_file: Path,
    max_size: int,
    max_backups: int,
) -> None:
    """Configure logging with rotating file handler and optional stdout."""
    root = logging.getLogger()
    root.setLevel(logging.DEBUG)

    fmt = logging.Formatter("%(asctime)s %(levelname)-5s %(message)s", datefmt="%Y-%m-%d %H:%M:%S")

    log_file.parent.mkdir(parents=True, exist_ok=True)
    file_handler = logging.handlers.RotatingFileHandler(
        log_file,
        maxBytes=max_size,
        backupCount=max_backups,
        encoding="utf-8",
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(fmt)
    root.addHandler(file_handler)

    if verbose:
        stdout_handler = logging.StreamHandler(sys.stdout)
        stdout_handler.setLevel(logging.DEBUG)
        stdout_handler.setFormatter(fmt)
        root.addHandler(stdout_handler)


# --- Main ---------------------------------------------------------------------


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        prog="roko-backup.py",
        description="Roko backup: collect -> archive -> upload -> verify -> clean -> report",
    )
    parser.add_argument(
        "--config", type=Path, default=None,
        help="Config file path (default: env vars / built-in defaults)",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Collect and print what would be backed up, don't archive/upload",
    )
    parser.add_argument(
        "--verbose", action="store_true",
        help="Detailed logging output (DEBUG level, also enables stdout)",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    """Entry point."""
    args = parse_args(argv)

    if args.config:
        config = Config.from_file(args.config)
    else:
        config = Config.from_env()

    setup_logging(
        verbose=args.verbose,
        log_file=config.log_file,
        max_size=config.max_log_size,
        max_backups=config.max_log_backups,
    )

    if not shutil.which("rclone"):
        log.error("rclone not found in PATH")
        return 1

    lock = LockFile(config.lock_path)
    try:
        with lock:
            pipeline = BackupPipeline(config, dry_run=args.dry_run)
            return pipeline.run()
    except LockError as exc:
        log.error(str(exc))
        return 2
    except Exception as exc:
        log.exception("Unexpected error: %s", exc)
        Reporter(config).write("error", error=str(exc))
        return 3


if __name__ == "__main__":
    sys.exit(main())