#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业15 · 生成「Dify 备份清单」（DB + 存储 + 向量相关）

扫描 docker-compose 与 volumes 目录，输出 Markdown/JSON 备份检查表。

用法：
  python hw15_01_备份清单.py
  python hw15_01_备份清单.py --dify-root /opt/dify-hw15
  python hw15_01_备份清单.py --vector-store weaviate --db-type postgresql

输出：
  data/homework15/dify_backup_checklist_YYYYMMDD.md
  data/homework15/dify_backup_checklist_YYYYMMDD.json
"""
from __future__ import annotations

import argparse
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Any


# 与官方 compose 默认挂载一致（本仓库 docker-compose.yaml）
VOLUME_ITEMS: list[dict[str, str]] = [
    {
        "id": "db_postgres",
        "category": "database",
        "path": "volumes/db/data",
        "desc": "PostgreSQL 主库（应用元数据、用户、应用配置、对话索引状态等）",
        "backup": "pg_dump 逻辑备份 + 卷目录快照（停写或一致性快照）",
        "restore_note": "恢复后需与 Redis/向量库版本一致；插件库 dify_plugin 在 PG 内",
    },
    {
        "id": "db_mysql",
        "category": "database",
        "path": "volumes/mysql/data",
        "desc": "MySQL 主库（DB_TYPE=mysql 时）",
        "backup": "mysqldump + 卷目录快照",
        "restore_note": "与 postgresql 二选一，勿混备",
    },
    {
        "id": "redis",
        "category": "cache_queue",
        "path": "volumes/redis/data",
        "desc": "Redis（Celery broker、缓存、Agent 会话态 db/2 等）",
        "backup": "可选：RDB/AOF 文件；生产更常接受「丢队列可重建」",
        "restore_note": "Worker 宕机实验后队列堆积在此；恢复 worker 可消费，不必长期备 Redis",
    },
    {
        "id": "app_storage",
        "category": "object_storage",
        "path": "volumes/app/storage",
        "desc": "用户上传文件、知识库原文、本地 STORAGE 模式对象",
        "backup": "目录 rsync/tar；若 STORAGE_TYPE=s3 则改备对象存储桶",
        "restore_note": "与 PG 中 documents 记录必须一致，否则出现「有记录无文件」",
    },
    {
        "id": "weaviate",
        "category": "vector",
        "path": "volumes/weaviate",
        "desc": "Weaviate 向量数据（VECTOR_STORE=weaviate）",
        "backup": "卷目录快照；大规模可用 Weaviate 备份 API",
        "restore_note": "与 PG 中 segment/document 索引状态一致；换库需全量重建索引",
    },
    {
        "id": "milvus_etcd",
        "category": "vector",
        "path": "volumes/milvus/etcd",
        "desc": "Milvus etcd 元数据（profile milvus）",
        "backup": "与 milvus/minio/milvus 三件套一起备",
        "restore_note": "仅 VECTOR_STORE=milvus 时启用",
    },
    {
        "id": "milvus_minio",
        "category": "vector",
        "path": "volumes/milvus/minio",
        "desc": "Milvus 对象块存储",
        "backup": "目录快照",
        "restore_note": "与 milvus standalone 配套",
    },
    {
        "id": "milvus_data",
        "category": "vector",
        "path": "volumes/milvus/milvus",
        "desc": "Milvus standalone 数据",
        "backup": "目录快照",
        "restore_note": "与 milvus profile 配套",
    },
    {
        "id": "pgvector",
        "category": "vector",
        "path": "volumes/pgvector/data",
        "desc": "pgvector（VECTOR_STORE=pgvector）",
        "backup": "pg_dump + 卷快照",
        "restore_note": "向量与关系数据同库",
    },
    {
        "id": "plugin_daemon",
        "category": "plugin",
        "path": "volumes/plugin_daemon",
        "desc": "插件包、插件运行时工作目录",
        "backup": "目录 rsync/tar",
        "restore_note": "另需 PG 库 dify_plugin",
    },
    {
        "id": "sandbox",
        "category": "aux",
        "path": "volumes/sandbox",
        "desc": "代码沙箱依赖与配置",
        "backup": "可选；丢后可重装依赖",
        "restore_note": "非核心业务数据",
    },
    {
        "id": "certbot",
        "category": "tls",
        "path": "volumes/certbot",
        "desc": "Let's Encrypt 证书",
        "backup": "conf/live 目录",
        "restore_note": "HTTPS 启用时需要",
    },
]

CELERY_QUEUES = [
    {"name": "dataset", "owner": "worker", "symptom_if_stuck": "知识库文档一直「索引中」"},
    {"name": "mail", "owner": "worker", "symptom_if_stuck": "邀请/通知邮件不发"},
    {"name": "ops_trace", "owner": "worker", "symptom_if_stuck": "Langfuse 等追踪延迟或缺失"},
    {"name": "app_deletion", "owner": "worker", "symptom_if_stuck": "删除应用/数据集卡住"},
    {"name": "workflow", "owner": "worker", "symptom_if_stuck": "部分异步工作流节点不推进"},
]

ENV_FILES = [
    ".env",
    "envs/core-services/shared.env",
    "envs/databases/db-postgres.env",
    "envs/databases/redis.env",
    "envs/vectorstores/weaviate.env",
    "nginx/nginx.conf.template",
]


def out_dir() -> Path:
    env = os.environ.get("HW15_OUT")
    if env:
        p = Path(env)
        p.mkdir(parents=True, exist_ok=True)
        return p
    repo = Path(__file__).resolve().parents[2] / "data" / "homework15"
    candidates = [repo]
    opt = Path("/opt/workbuddy/data/homework15")
    if os.name != "nt" and Path("/opt/workbuddy").is_dir():
        candidates.insert(0, opt)
    for c in candidates:
        try:
            c.mkdir(parents=True, exist_ok=True)
            return c
        except OSError:
            continue
    return repo


def dir_size_human(path: Path) -> str:
    if not path.exists():
        return "不存在"
    total = 0
    try:
        for root, _dirs, files in os.walk(path):
            for f in files:
                fp = Path(root) / f
                try:
                    total += fp.stat().st_size
                except OSError:
                    pass
    except OSError:
        return "不可读"
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if total < 1024:
            return f"{total:.1f} {unit}"
        total /= 1024
    return f"{total:.1f} PB"


def scan_volumes(root: Path, vector_store: str, db_type: str) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for item in VOLUME_ITEMS:
        cat = item["category"]
        vid = item["id"]
        if cat == "database" and db_type == "postgresql" and vid == "db_mysql":
            continue
        if cat == "database" and db_type == "mysql" and vid == "db_postgres":
            continue
        if cat == "vector":
            if vector_store == "weaviate" and vid.startswith("milvus"):
                continue
            if vector_store == "milvus" and vid in ("weaviate", "pgvector"):
                continue
            if vector_store == "pgvector" and vid not in ("pgvector",):
                if vid.startswith("milvus") or vid == "weaviate":
                    continue
            if vector_store not in ("weaviate", "milvus", "pgvector") and vid in (
                "weaviate",
                "milvus_etcd",
                "milvus_minio",
                "milvus_data",
                "pgvector",
            ):
                if vector_store != "weaviate" and vid == "weaviate":
                    continue
        abs_path = root / item["path"]
        rows.append(
            {
                **item,
                "abs_path": str(abs_path),
                "exists": abs_path.exists(),
                "size": dir_size_human(abs_path),
                "required": _is_required(item, vector_store, db_type),
            }
        )
    return rows


def _is_required(item: dict[str, str], vector_store: str, db_type: str) -> bool:
    if item["category"] == "database":
        if db_type == "postgresql":
            return item["id"] == "db_postgres"
        return item["id"] == "db_mysql"
    if item["category"] == "object_storage":
        return True
    if item["category"] == "vector":
        if vector_store == "weaviate":
            return item["id"] == "weaviate"
        if vector_store == "milvus":
            return item["id"].startswith("milvus")
        if vector_store == "pgvector":
            return item["id"] == "pgvector"
    if item["id"] in ("redis", "plugin_daemon"):
        return True
    return False


def build_checklist(root: Path, vector_store: str, db_type: str) -> dict[str, Any]:
    volumes = scan_volumes(root, vector_store, db_type)
    env_status = []
    for rel in ENV_FILES:
        p = root / rel
        env_status.append({"path": rel, "exists": p.exists()})

    return {
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "dify_root": str(root),
        "vector_store": vector_store,
        "db_type": db_type,
        "volumes": volumes,
        "celery_queues": CELERY_QUEUES,
        "env_files": env_status,
        "backup_order": [
            "1. 记录当前镜像 tag 与 .env（版本一致性）",
            "2. 停写或低峰：pg_dump / mysqldump 主库 + dify_plugin",
            "3. rsync/tar volumes/app/storage（或 S3 桶清单）",
            "4. 向量卷快照（weaviate / milvus / pgvector 按 VECTOR_STORE）",
            "5. volumes/plugin_daemon",
            "6. （可选）Redis；队列类数据通常可重建",
            "7. 校验：恢复演练环境导入后 spot-check 知识库检索 + 对话",
        ],
    }


def to_markdown(data: dict[str, Any]) -> str:
    lines = [
        "# Dify 备份清单（自动生成）",
        "",
        f"- 生成时间：`{data['generated_at']}`",
        f"- Dify 根目录：`{data['dify_root']}`",
        f"- 向量库：`{data['vector_store']}` · 主库：`{data['db_type']}`",
        "",
        "## 1. 必备数据卷",
        "",
        "| 必備 | 类别 | 路径 | 大小 | 说明 | 备份方式 |",
        "|------|------|------|------|------|----------|",
    ]
    for v in data["volumes"]:
        req = "是" if v["required"] else "否"
        lines.append(
            f"| {req} | {v['category']} | `{v['path']}` | {v['size']} | {v['desc']} | {v['backup']} |"
        )

    lines.extend(
        [
            "",
            "## 2. Celery 队列（Worker 相关，Redis broker）",
            "",
            "| 队列 | 消费者 | Worker 停掉时的用户现象 |",
            "|------|--------|-------------------------|",
        ]
    )
    for q in data["celery_queues"]:
        lines.append(f"| `{q['name']}` | {q['owner']} | {q['symptom_if_stuck']} |")

    lines.extend(["", "## 3. 配置文件", ""])
    for e in data["env_files"]:
        mark = "✓" if e["exists"] else "✗"
        lines.append(f"- [{mark}] `{e['path']}`")

    lines.extend(["", "## 4. 推荐备份顺序", ""])
    for step in data["backup_order"]:
        lines.append(f"- {step}")

    lines.extend(
        [
            "",
            "## 5. 恢复注意",
            "",
            "- PostgreSQL/MySQL 与 `volumes/app/storage`、向量库 **三者时间点尽量一致**",
            "- 仅恢复 PG 不恢复向量 → 检索空或报错；仅恢复向量不恢复 PG → 孤儿向量",
            "- `STORAGE_TYPE` 非 local 时，对象存储桶与 PG 记录同等重要",
            "- 插件：`dify_plugin` 库 + `volumes/plugin_daemon` 一起恢复",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description="生成 Dify 备份清单")
    parser.add_argument(
        "--dify-root",
        default=None,
        help="Dify docker-compose 根目录（默认：仓库根或 /opt/workbuddy）",
    )
    parser.add_argument("--vector-store", default="weaviate", choices=["weaviate", "milvus", "pgvector"])
    parser.add_argument("--db-type", default="postgresql", choices=["postgresql", "mysql"])
    args = parser.parse_args()

    if args.dify_root:
        root = Path(args.dify_root)
    else:
        repo = Path(__file__).resolve().parents[2]
        root = Path("/opt/workbuddy") if (os.name != "nt" and Path("/opt/workbuddy").is_dir()) else repo

    data = build_checklist(root, args.vector_store, args.db_type)
    od = out_dir()
    stamp = datetime.now().strftime("%Y%m%d")
    md_path = od / f"dify_backup_checklist_{stamp}.md"
    json_path = od / f"dify_backup_checklist_{stamp}.json"
    md_path.write_text(to_markdown(data), encoding="utf-8")
    json_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Wrote {md_path}")
    print(f"Wrote {json_path}")
    req = [v for v in data["volumes"] if v["required"]]
    print(f"\n必備卷 {len(req)} 项：")
    for v in req:
        print(f"  - {v['path']} ({v['size']})")


if __name__ == "__main__":
    main()