#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# backup.py —— 操作前自动留备份，支持查看与恢复

"""
每次写文件、删文件/目录、执行 shell/Python 之前，把旧内容存一份：

  user/backups/<备份 id>/
      meta.json        元信息：类型、原路径、时间、工具、大小、哈希
      payload          单文件备份（原样字节）
      payload.tar.gz   目录 / 工作区 / 记忆备份

备份 id 形如 20261105-113500-a1b2c3，按时间排序。保留策略按条数和总大小裁剪，
最新的和正在恢复的那一份不会被删。

执行前只给记忆文件（user/*.MD）拍快照：工作目录可能几十万个文件，整目录打包太贵。
记忆快照按“文件名+大小+mtime”做签名，内容没变就跳过。
"""

import hashlib
import json
import logging
import os
import re
import shutil
import tarfile
import threading
import time
import uuid
from datetime import datetime
from pathlib import Path

from scheduler import resolve_tz

logger = logging.getLogger("AgentBackup")

BACKUP_ID_RE = re.compile(r"^\d{8}-\d{6}-[0-9a-f]{6}$")
MAX_FILE_BYTES_HARD = 1024 * 1024 * 1024      # 单文件硬上限 1GB，防止误配置


class BackupError(Exception):
    pass


class BackupTooLarge(BackupError):
    pass


class BackupTimeout(BackupError):
    """遍历/打包超过时间预算；宁可跳过快照，也不要把一条命令堵几分钟。"""


def _env_int(name, default):
    try:
        return int(os.getenv(name, str(default)))
    except (TypeError, ValueError):
        return default


def enabled():
    return os.getenv("AGENT_BACKUP_ENABLED", "1").strip().lower() not in (
        "0", "false", "no", "off")


def before_exec_enabled():
    return os.getenv("AGENT_BACKUP_BEFORE_EXEC", "1").strip().lower() not in (
        "0", "false", "no", "off")


def default_backup_root():
    env_path = os.getenv("AGENT_BACKUP_DIR", "").strip()
    if env_path:
        return Path(env_path)
    return Path(__file__).resolve().parent / "user" / "backups"


def _sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _is_within(path, parent):
    try:
        path.relative_to(parent)
        return True
    except ValueError:
        return False


def _iter_tree(root, exclude=None, skip_logs=False, deadline=None):
    """稳定顺序逐个产出条目；超过 deadline 抛 BackupTimeout。

    用 os.walk 惰性遍历：rglob + sorted 会先把整棵树物化成列表，工作目录一挂到
    Docker Desktop 的大目录（node_modules 之类）上，光这一步就能堵几分钟。
    """
    root = Path(root)

    def _skipped(p):
        if exclude is not None and (p == exclude or _is_within(p, exclude)):
            return True
        return bool(skip_logs and p.name.startswith("audit.log"))

    for dirpath, dirnames, filenames in os.walk(root, topdown=True,
                                                followlinks=False):
        if deadline is not None and time.monotonic() > deadline:
            raise BackupTimeout("遍历工作目录超过时间预算")
        base = Path(dirpath)
        dirnames.sort()
        filenames.sort()
        dirnames[:] = [d for d in dirnames if not _skipped(base / d)]
        for name in dirnames:
            yield base / name
        for name in filenames:
            p = base / name
            if _skipped(p):
                continue
            if deadline is not None and time.monotonic() > deadline:
                raise BackupTimeout("遍历工作目录超过时间预算")
            yield p


def tree_signature(root, exclude=None, skip_logs=True, deadline=None):
    """目录状态签名：路径 + 大小 + mtime；备份目录和日志不参与。"""
    root = Path(root)
    digest = hashlib.sha256()
    for p in _iter_tree(root, exclude=exclude, skip_logs=skip_logs,
                        deadline=deadline):
        try:
            rel = p.relative_to(root).as_posix()
            if p.is_symlink():
                digest.update(("L " + rel + " " + os.readlink(p) + "\n").encode())
            elif p.is_dir():
                digest.update(("D " + rel + "\n").encode())
            else:
                st = p.stat()
                digest.update(
                    ("F " + rel + " " + str(st.st_size) + " "
                     + str(st.st_mtime_ns) + "\n").encode())
        except OSError:
            continue
    return digest.hexdigest()


class BackupStore:
    def __init__(self, root=None, max_items=None, max_total_mb=None,
                 file_max_mb=None, tree_max_mb=None, workspace_max_mb=None):
        self.root = Path(root or default_backup_root()).resolve()
        self.max_items = max(1, max_items if max_items is not None
                             else _env_int("AGENT_BACKUP_MAX_ITEMS", 200))
        self.max_total_bytes = max(1, max_total_mb if max_total_mb is not None
                                   else _env_int("AGENT_BACKUP_MAX_MB", 500)) * 1024 * 1024
        self.file_max_bytes = min(
            MAX_FILE_BYTES_HARD,
            max(1, file_max_mb if file_max_mb is not None
                else _env_int("AGENT_BACKUP_FILE_MAX_MB", 50)) * 1024 * 1024)
        self.tree_max_bytes = max(1, tree_max_mb if tree_max_mb is not None
                                  else _env_int("AGENT_BACKUP_TREE_MAX_MB", 100)) * 1024 * 1024
        self.workspace_max_bytes = max(
            1, workspace_max_mb if workspace_max_mb is not None
            else _env_int("AGENT_BACKUP_WORKSPACE_MB", 50)) * 1024 * 1024
        self._lock = threading.RLock()
        self._last_seq = 0
        self._last_workspace = None      # (signature, backup_id)
        self._last_memory = None         # (signature, backup_id)
        self.root.mkdir(parents=True, exist_ok=True)

    # ---------------- 内部工具 ----------------
    def _new_id(self):
        return datetime.now(resolve_tz()).strftime("%Y%m%d-%H%M%S") + "-" + uuid.uuid4().hex[:6]

    def _next_seq(self):
        """同一秒内多次备份也能稳定排序。"""
        seq = time.time_ns()
        if seq <= self._last_seq:
            seq = self._last_seq + 1
        self._last_seq = seq
        return seq

    def _dir(self, backup_id):
        if not BACKUP_ID_RE.match(str(backup_id or "")):
            raise BackupError(f"备份 id 格式不对：{backup_id}")
        return self.root / backup_id

    def _write_meta(self, bdir, meta):
        (bdir / "meta.json").write_text(
            json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")

    def _read_meta(self, bdir):
        try:
            meta = json.loads((bdir / "meta.json").read_text(encoding="utf-8"))
        except (OSError, ValueError):
            return None
        if not isinstance(meta, dict) or not meta.get("id"):
            return None
        meta["_dir"] = str(bdir)
        return meta

    def _iter_metas(self):
        metas = []
        try:
            children = list(self.root.iterdir())
        except OSError:
            return metas
        for child in children:
            if not child.is_dir() or not BACKUP_ID_RE.match(child.name):
                continue
            meta = self._read_meta(child)
            if meta:
                metas.append(meta)
        metas.sort(key=lambda m: (m.get("seq") or 0, m.get("ts") or "",
                                  m.get("id") or ""), reverse=True)
        return metas

    def _measure(self, path, exclude=None, skip_logs=False, deadline=None,
                 limit=None):
        """返回 (总字节数, 文件数)；limit 给定时一超就返回，不再白走整棵树。"""
        path = Path(path)
        total = 0
        count = 0
        if path.is_file():
            return path.stat().st_size, 1
        for p in _iter_tree(path, exclude=exclude, skip_logs=skip_logs,
                            deadline=deadline):
            try:
                if p.is_file() and not p.is_symlink():
                    total += p.stat().st_size
                    count += 1
                    if limit is not None and total > limit:
                        return total, count
            except OSError:
                continue
        return total, count

    def _prune(self, keep_ids=None):
        """按条数和总大小裁剪，从最旧的开始删。"""
        keep_ids = set(keep_ids or ())
        metas = self._iter_metas()
        total = sum(int(m.get("stored_bytes") or 0) for m in metas)
        for meta in reversed(list(metas)):
            if len(metas) <= self.max_items and total <= self.max_total_bytes:
                break
            if meta["id"] in keep_ids:
                continue
            bdir = Path(meta["_dir"])
            try:
                shutil.rmtree(bdir)
                total -= int(meta.get("stored_bytes") or 0)
                metas.remove(meta)
                logger.info(f"清理旧备份 {meta['id']}")
            except OSError as e:
                logger.warning(f"清理旧备份失败（{meta['id']}）：{e}")

    # ---------------- 创建备份 ----------------
    def create_file(self, path, tool="", action="", note="", keep_ids=None):
        path = Path(path)
        if path == self.root or _is_within(path, self.root):
            raise BackupError("不能把备份目录本身再备份一遍")
        if not path.is_file():
            return None
        size = path.stat().st_size
        if size > self.file_max_bytes:
            raise BackupTooLarge(
                f"文件 {path} 有 {size} 字节，超过单文件备份上限 "
                f"{self.file_max_bytes // (1024 * 1024)}MB")

        with self._lock:
            backup_id = self._new_id()
            bdir = self.root / backup_id
            bdir.mkdir(mode=0o700)
            try:
                payload = bdir / "payload"
                shutil.copy2(path, payload)
                meta = {
                    "id": backup_id,
                    "seq": self._next_seq(),
                    "ts": datetime.now(resolve_tz()).isoformat(timespec="seconds"),
                    "kind": "file",
                    "tool": tool,
                    "action": action,
                    "original": str(path),
                    "size": size,
                    "stored_bytes": payload.stat().st_size,
                    "sha256": _sha256(payload),
                    "note": str(note or "")[:300],
                }
                self._write_meta(bdir, meta)
            except Exception:
                shutil.rmtree(bdir, ignore_errors=True)
                raise

            self._prune(keep_ids={backup_id} | set(keep_ids or ()))
            logger.info(f"已备份文件 {path} → {backup_id}")
            return dict(meta)

    def create_tree(self, path, tool="", action="", kind="directory", note="",
                    signature=None, keep_ids=None, skip_logs=False,
                    deadline=None):
        path = Path(path)
        if path == self.root or _is_within(path, self.root):
            raise BackupError("不能把备份目录本身再备份一遍")
        if not path.is_dir():
            return None
        limit = self.workspace_max_bytes if kind == "workspace" else self.tree_max_bytes
        total, count = self._measure(path, exclude=self.root,
                                     skip_logs=skip_logs, deadline=deadline,
                                     limit=limit)
        if total > limit:
            raise BackupTooLarge(
                f"目录 {path} 超过上限 {limit // (1024 * 1024)}MB"
                f"（已扫描 {total // (1024 * 1024)}MB），未备份")

        with self._lock:
            backup_id = self._new_id()
            bdir = self.root / backup_id
            bdir.mkdir(mode=0o700)
            try:
                payload = bdir / "payload.tar.gz"
                self._tar_tree(path, payload, skip_logs=skip_logs,
                               deadline=deadline)
                meta = {
                    "id": backup_id,
                    "seq": self._next_seq(),
                    "ts": datetime.now(resolve_tz()).isoformat(timespec="seconds"),
                    "kind": kind,
                    "tool": tool,
                    "action": action,
                    "original": str(path),
                    "size": total,
                    "files": count,
                    "stored_bytes": payload.stat().st_size,
                    "sha256": _sha256(payload),
                    "signature": signature,
                    "note": str(note or "")[:300],
                }
                self._write_meta(bdir, meta)
            except Exception:
                shutil.rmtree(bdir, ignore_errors=True)
                raise

            self._prune(keep_ids={backup_id} | set(keep_ids or ()))
            logger.info(f"已备份目录 {path} → {backup_id}（{count} 个文件）")
            return dict(meta)

    def _tar_tree(self, root, target, skip_logs=False, deadline=None):
        root = Path(root)
        with tarfile.open(target, "w:gz") as tar:
            tar.add(root, arcname=".", recursive=False)
            for p in _iter_tree(root, exclude=self.root, skip_logs=skip_logs,
                                deadline=deadline):
                tar.add(p, arcname=p.relative_to(root).as_posix(),
                        recursive=False)

    def create_workspace(self, path, tool="", note="", deadline=None):
        """执行前快照；目录没变化时返回 None（跳过）。"""
        path = Path(path).resolve()
        total, count = self._measure(path, exclude=self.root, skip_logs=True,
                                     deadline=deadline,
                                     limit=self.workspace_max_bytes)
        if total > self.workspace_max_bytes:
            raise BackupTooLarge(
                f"工作目录 {path} 超过快照上限 "
                f"{self.workspace_max_bytes // (1024 * 1024)}MB"
                f"（已扫描 {total // (1024 * 1024)}MB）")

        signature = tree_signature(path, exclude=self.root, skip_logs=True,
                                   deadline=deadline)
        with self._lock:
            if self._last_workspace is None:
                for meta in self._iter_metas():
                    if meta.get("kind") == "workspace" and meta.get("signature"):
                        self._last_workspace = (meta["signature"], meta["id"])
                        break
            last_sig, last_id = self._last_workspace or (None, None)
            if signature == last_sig and last_id and (self.root / last_id).exists():
                return None

            meta = self.create_tree(
                path, tool=tool, action="exec", kind="workspace",
                note=note, signature=signature, skip_logs=True,
                deadline=deadline)
            if meta:
                self._last_workspace = (signature, meta["id"])
            return meta

    def create_memory(self, directory, tool="", note="", deadline=None):
        """记忆快照：只打包目录下的 *.MD，工作目录再大也不碰。

        USER.MD / MEMORY.MD / MEMORY-日期.MD 是长期资产，体积小、改动频繁；
        每次执行前给它们留一份，比整目录快照便宜得多。内容没变化时返回 None。
        """
        directory = Path(directory)
        try:
            files = sorted(
                (p for p in directory.iterdir()
                 if p.is_file() and p.suffix.lower() == ".md"),
                key=lambda p: p.name)
        except OSError:
            return None
        if not files:
            return None

        digest = hashlib.sha256()
        total = 0
        for p in files:
            if deadline is not None and time.monotonic() > deadline:
                raise BackupTimeout("记忆快照超过时间预算")
            try:
                st = p.stat()
            except OSError:
                continue
            total += st.st_size
            digest.update(
                ("F " + p.name + " " + str(st.st_size) + " "
                 + str(st.st_mtime_ns) + "\n").encode())
        signature = digest.hexdigest()

        if total > self.workspace_max_bytes:
            raise BackupTooLarge(
                f"记忆目录 {directory} 超过快照上限 "
                f"{self.workspace_max_bytes // (1024 * 1024)}MB")

        with self._lock:
            if self._last_memory is None:
                for meta in self._iter_metas():
                    if meta.get("kind") == "memory" and meta.get("signature"):
                        self._last_memory = (meta["signature"], meta["id"])
                        break
            last_sig, last_id = self._last_memory or (None, None)
            if signature == last_sig and last_id and (self.root / last_id).exists():
                return None

            backup_id = self._new_id()
            bdir = self.root / backup_id
            bdir.mkdir(mode=0o700)
            try:
                payload = bdir / "payload.tar.gz"
                with tarfile.open(payload, "w:gz") as tar:
                    for p in files:
                        if deadline is not None and time.monotonic() > deadline:
                            raise BackupTimeout("记忆快照超过时间预算")
                        try:
                            tar.add(p, arcname=p.name, recursive=False)
                        except OSError:
                            continue
                meta = {
                    "id": backup_id,
                    "seq": self._next_seq(),
                    "ts": datetime.now(resolve_tz()).isoformat(timespec="seconds"),
                    "kind": "memory",
                    "tool": tool,
                    "action": "exec",
                    "original": str(directory),
                    "size": total,
                    "files": len(files),
                    "stored_bytes": payload.stat().st_size,
                    "sha256": _sha256(payload),
                    "signature": signature,
                    "note": str(note or "")[:300],
                }
                self._write_meta(bdir, meta)
            except Exception:
                shutil.rmtree(bdir, ignore_errors=True)
                raise

            self._prune(keep_ids={backup_id})
            self._last_memory = (signature, backup_id)
            logger.info(f"已备份记忆 {directory} → {backup_id}（{len(files)} 个文件）")
            return dict(meta)

    # ---------------- 查询 ----------------
    def list(self, limit=20, path=None):
        try:
            limit = max(1, int(limit))
        except (TypeError, ValueError):
            limit = 20
        filter_path = str(Path(path).resolve()) if path else None
        metas = []
        for meta in self._iter_metas():
            if filter_path and meta.get("original") != filter_path:
                continue
            metas.append({k: v for k, v in meta.items() if k != "_dir"})
            if len(metas) >= limit:
                break
        return metas

    def get(self, backup_id):
        meta = self._read_meta(self._dir(backup_id))
        if meta:
            meta.pop("_dir", None)
        return meta

    # ---------------- 恢复 ----------------
    def restore(self, backup_id, target=None):
        meta = self.get(backup_id)
        if not meta:
            raise BackupError(f"没有这个备份：{backup_id}")

        bdir = self._dir(backup_id)
        target_path = Path(target).resolve() if target else Path(meta["original"]).resolve()
        if target_path == self.root or _is_within(target_path, self.root):
            raise BackupError("不能把备份恢复到备份目录里")

        kind = meta.get("kind")
        if kind == "file":
            payload = bdir / "payload"
            if not payload.exists():
                raise BackupError("备份文件缺失")
            if meta.get("sha256") and _sha256(payload) != meta["sha256"]:
                raise BackupError("备份文件校验失败，可能已损坏")
            if target_path.exists():
                if target_path.is_dir():
                    raise BackupError(f"目标是一个目录：{target_path}")
                self.create_file(target_path, tool="restore_backup",
                                 action="restore-overwrite", keep_ids={backup_id})
            target_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(payload, target_path)

        elif kind in ("directory", "workspace", "memory"):
            payload = bdir / "payload.tar.gz"
            if not payload.exists():
                raise BackupError("备份压缩包缺失")
            if meta.get("sha256") and _sha256(payload) != meta["sha256"]:
                raise BackupError("备份压缩包校验失败，可能已损坏")

            if target_path.exists():
                if target_path.is_file():
                    raise BackupError(f"目标是一个文件：{target_path}")
                # 恢复前先把当前目录也备份一份，恢复本身也可回退
                self.create_tree(target_path, tool="restore_backup",
                                 action="restore-overwrite", kind="directory",
                                 keep_ids={backup_id})
            else:
                target_path.mkdir(parents=True, exist_ok=True)

            with tarfile.open(payload, "r:gz") as tar:
                try:
                    tar.extractall(target_path, filter="data")
                except TypeError:
                    # 老版本 Python 没有 filter 参数
                    tar.extractall(target_path)
        else:
            raise BackupError(f"未知的备份类型：{kind}")

        logger.info(f"已从 {backup_id} 恢复到 {target_path}")
        result = dict(meta)
        result["restored_to"] = str(target_path)
        return result

    def cleanup(self):
        with self._lock:
            self._prune()


# ---------------- 进程内单例 ----------------
_STORE = None
_STORE_LOCK = threading.Lock()


def get_store():
    global _STORE
    with _STORE_LOCK:
        if _STORE is None:
            _STORE = BackupStore()
        return _STORE


def set_store(store):
    global _STORE
    with _STORE_LOCK:
        _STORE = store
    return _STORE