"""
把 Artalk 导出的 .artrans 转成可直接灌进 Cloudflare D1 的 SQL。

为什么不用 Worker 的 /transfer/import：
  3440 条评论逐条 INSERT 会产生 3440 次 D1 往返，Worker 免费版单请求 CPU 只有 10ms，
  很容易直接超时。批量导入走 wrangler d1 execute --file 最稳、最快。

用法：
  python tools/import_artrans.py                    # 默认只导「优世界」
  python tools/import_artrans.py --site 优世界 --src <path.artrans> --out out/
  # 然后：
  #   npx wrangler d1 execute artalk-cf --remote --file=out/import-01-users.sql
  #   ... 按生成的 run.sh 顺序执行
"""

import argparse
import json
import pathlib
import sys
from collections import OrderedDict

DEFAULT_SRC = ""  # 不写死机器路径；用 --src <path.artrans> 显式传入
DEFAULT_SITE = "优世界"
DEFAULT_SITE_URL = "https://usj.cc"

CHUNK = 400  # 每个 SQL 文件最多多少条语句，避免单次 execute 过大


def sql_str(v) -> str:
    if v is None:
        return "NULL"
    s = str(v)
    return "'" + s.replace("'", "''") + "'"


def sql_bool(v) -> str:
    return "1" if str(v).lower() in ("true", "1", "yes") else "0"


def sql_int(v, dft=0) -> str:
    try:
        return str(int(str(v).strip()))
    except Exception:
        return str(dft)


def parse_dt(s: str) -> int:
    """'2021-09-11 20:33:56 +0800 +0800' -> unix 毫秒（+0800）"""
    s = (s or "").strip()
    if len(s) < 19:
        return 0
    try:
        import datetime as dt

        naive = dt.datetime.strptime(s[:19], "%Y-%m-%d %H:%M:%S")
        aware = naive.replace(tzinfo=dt.timezone(dt.timedelta(hours=8)))
        ms = int(aware.timestamp() * 1000)
        return ms if ms > 0 else 0
    except Exception:
        return 0


def write_chunks(out_dir: pathlib.Path, name: str, stmts: list) -> list:
    """把语句按 CHUNK 切成多个 .sql，返回生成的文件名列表"""
    files = []
    for i in range(0, len(stmts), CHUNK):
        part = stmts[i : i + CHUNK]
        fn = f"{name}-{i // CHUNK + 1:03d}.sql"
        (out_dir / fn).write_text("\n".join(part) + "\n", encoding="utf-8", newline="\n")
        files.append(fn)
    return files


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=DEFAULT_SRC)
    ap.add_argument("--site", default=DEFAULT_SITE, help="只导这个站点名，逗号分隔可多个；all = 全导")
    ap.add_argument("--site-url", default=DEFAULT_SITE_URL)
    ap.add_argument("--out", default="tools/import-out")
    ap.add_argument("--db", default="artalk-cf", help="wrangler d1 的数据库名")
    args = ap.parse_args()

    src = pathlib.Path(args.src)
    if not args.src:
        print("请用 --src <path.artrans> 指定导出的 .artrans 文件路径", file=sys.stderr)
        return 1
    if not src.exists():
        print(f"找不到源文件: {src}", file=sys.stderr)
        return 1

    data = json.loads(src.read_text(encoding="utf-8"))
    want = None if args.site.strip().lower() == "all" else [s.strip() for s in args.site.split(",") if s.strip()]
    site_name = want[0] if want else DEFAULT_SITE

    items = [c for c in data if (want is None or c.get("site_name") in want)]
    print(f"源文件 {len(data)} 条，按站点筛选后 {len(items)} 条 -> 归入站点「{site_name}」")
    if not items:
        print("筛选后为空，检查 --site", file=sys.stderr)
        return 1

    out_dir = pathlib.Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ---------------- 站点
    head = [
        "-- =====================================================================",
        "--  ⚠️ 这个文件会先清空 comments / pages / users / sites 四张表。",
        "--     只在**首次导入到全新数据库**时执行；库里已有数据请勿再跑。",
        "-- =====================================================================",
        "DELETE FROM comments;",
        "DELETE FROM pages;",
        "DELETE FROM users;",
        "DELETE FROM sites;",
        f"INSERT INTO sites (id, name, urls, created_at, updated_at) "
        f"VALUES (1, {sql_str(site_name)}, {sql_str(args.site_url)}, "
        f"{parse_dt(max((c.get('created_at') or '') for c in items)) or 0}, 0);",
    ]

    # ---------------- 用户（按 nick+email 去重）
    users = OrderedDict()
    for c in items:
        nick = (c.get("nick") or "").strip()
        email = (c.get("email") or "").strip()
        if not nick or not email:
            continue
        users.setdefault(f"{nick}\u0000{email}", {"nick": nick, "email": email, "link": c.get("link") or "",
                                                   "badge_name": c.get("badge_name") or "",
                                                   "badge_color": c.get("badge_color") or ""})

    user_id = {}
    user_stmts = ["-- 用户（原数据没有注册账号，全部按 昵称+邮箱 去重重建）"]
    for idx, (key, u) in enumerate(users.items(), start=1):
        user_id[key] = idx
        ts = next((parse_dt(c.get("created_at")) for c in items
                   if (c.get("nick") or "").strip() == u["nick"] and (c.get("email") or "").strip() == u["email"]), 0)
        user_stmts.append(
            "INSERT INTO users (id, name, email, link, badge_name, badge_color, receive_email, created_at, updated_at) "
            f"VALUES ({idx}, {sql_str(u['nick'])}, {sql_str(u['email'])}, {sql_str(u['link'])}, "
            f"{sql_str(u['badge_name'])}, {sql_str(u['badge_color'])}, 1, {ts or 0}, {ts or 0});"
        )
    print(f"用户 {len(users)} 个")

    # ---------------- 页面
    pages = OrderedDict()
    for c in items:
        key = (c.get("page_key") or "").strip()
        if not key:
            continue
        pages.setdefault(key, {"title": c.get("page_title") or "", "admin_only": c.get("page_admin_only"),
                               "created_at": parse_dt(c.get("created_at"))})

    page_id = {}
    page_stmts = ["-- 页面"]
    for idx, (key, p) in enumerate(pages.items(), start=1):
        page_id[key] = idx
        page_stmts.append(
            "INSERT INTO pages (id, key, site_name, title, admin_only, created_at, updated_at) "
            f"VALUES ({idx}, {sql_str(key)}, {sql_str(site_name)}, {sql_str(p['title'])}, "
            f"{sql_bool(p['admin_only'])}, {p['created_at'] or 0}, {p['created_at'] or 0});"
        )
    print(f"页面 {len(pages)} 个")

    # ---------------- 评论（保留原始 id，保证 rid 关系不变）
    rid_of = {}
    rows = []
    skip_no_id = skip_no_page = skip_no_user = 0
    for c in items:
        try:
            cid = int(str(c.get("id")))
        except Exception:
            skip_no_id += 1
            continue
        try:
            rid = int(str(c.get("rid") or 0))
        except Exception:
            rid = 0
        key = f"{(c.get('nick') or '').strip()}\u0000{(c.get('email') or '').strip()}"
        pk = (c.get("page_key") or "").strip()
        if not pk:
            skip_no_page += 1
            continue
        if key not in user_id:
            skip_no_user += 1
            continue
        rid_of[cid] = rid
        rows.append({
            "id": cid, "rid": rid, "uid": user_id[key], "pid": page_id[pk], "page_key": pk,
            "content": c.get("content") or "",
            "is_collapsed": c.get("is_collapsed"), "is_pending": c.get("is_pending"),
            "is_pinned": c.get("is_pinned"), "vote_up": c.get("vote_up"), "vote_down": c.get("vote_down"),
            "ua": c.get("ua") or "", "ip": c.get("ip") or "",
            "created_at": parse_dt(c.get("created_at")), "updated_at": parse_dt(c.get("updated_at")),
        })

    # 只保留指向本批次内评论的 rid，避免指到没导进来的兄弟站点评论
    ids = {r["id"] for r in rows}
    for r in rows:
        if r["rid"] and r["rid"] not in ids:
            r["rid"] = 0

    # 计算 root（沿 rid 链上溯到「rid==0 的那条顶层评论」）
    #
    # ★ 这里必须返回顶层评论自身的 id，而不是 0。
    #   对应官方 dao.FindCommentRootID：
    #     for rootId != 0 { if comment.Rid == 0 { return rootId }; rootId = comment.Rid }
    #   即「父链上第一个 rid 为 0 的那条评论」。
    #   注意语义是：顶层评论 root_id=0，它的所有子孙 root_id=该顶层评论的 id。
    root_cache = {}

    def root_of(cid: int) -> int:
        if cid in root_cache:
            return root_cache[cid]
        chain = []
        cur = cid
        seen = set()
        while cur and cur not in root_cache and cur not in seen:
            seen.add(cur)
            parent = rid_of.get(cur, 0)
            if parent == 0:
                break  # cur 就是顶层评论
            chain.append(cur)
            cur = parent
        root = root_cache.get(cur, cur)
        for n in chain:
            root_cache[n] = root
        root_cache[root] = root
        return root

    # 先全部插入成顶层，再统一回填 rid / root_id —— 与导入顺序无关
    ins, fix = [], []
    for r in rows:
        # is_verified 在 .artrans 里没有这个字段（官方导出不带）。
        # 前端只在「没有徽章」时才把 is_verified 渲染成小勾，所以这里统一给 0：
        #   有徽章的 328 条（博主）照样显示徽章，其余 3000+ 匿名评论不会多出勾。
        ins.append(
            "INSERT INTO comments (id, content, page_key, site_name, user_id, rid, root_id, "
            "is_collapsed, is_pending, is_pinned, is_verified, vote_up, vote_down, ua, ip, created_at, updated_at) "
            f"VALUES ({r['id']}, {sql_str(r['content'])}, {sql_str(r['page_key'])}, {sql_str(site_name)}, "
            f"{r['uid']}, 0, 0, {sql_bool(r['is_collapsed'])}, {sql_bool(r['is_pending'])}, {sql_bool(r['is_pinned'])}, "
            f"0, {sql_int(r['vote_up'])}, {sql_int(r['vote_down'])}, {sql_str(r['ua'])}, {sql_str(r['ip'])}, "
            f"{r['created_at']}, {r['updated_at'] or r['created_at']});"
        )
        if r["rid"]:
            fix.append(f"UPDATE comments SET rid = {r['rid']}, root_id = {root_of(r['rid'])} WHERE id = {r['id']};")
    print(f"评论 {len(rows)} 条（其中回复 {len(fix)} 条）")
    if skip_no_page or skip_no_user or skip_no_id:
        print(
            f"跳过 {skip_no_page + skip_no_user + skip_no_id} 条："
            f"page_key 为空 {skip_no_page}、昵称或邮箱为空 {skip_no_user}、id 非法 {skip_no_id}"
        )
        if skip_no_page:
            print("  ↑ page_key 为空的评论在官方 Artalk 里同样无法挂到任何页面，属于历史脏数据")

    # ---------------- 落盘
    files = []
    (out_dir / "00-reset-and-site.sql").write_text("\n".join(head) + "\n", encoding="utf-8", newline="\n")
    files.append("00-reset-and-site.sql")

    files += write_chunks(out_dir, "01-users", user_stmts)
    files += write_chunks(out_dir, "02-pages", page_stmts)
    files += write_chunks(out_dir, "03-comments", ins)
    files += write_chunks(out_dir, "04-rid-fix", fix)

    # 为了让 wrangler d1 execute 的 SQLite 序列号往后走
    tail = ["-- 让 AUTOINCREMENT 从最大值之后继续", f"UPDATE sqlite_sequence SET seq = "
            f"(SELECT MAX(id) FROM comments) WHERE name = 'comments';",
            f"UPDATE sqlite_sequence SET seq = (SELECT MAX(id) FROM users) WHERE name = 'users';",
            f"UPDATE sqlite_sequence SET seq = (SELECT MAX(id) FROM pages) WHERE name = 'pages';",
            "-- 自检",
            "SELECT 'sites' AS t, COUNT(*) AS n FROM sites",
            "UNION ALL SELECT 'pages', COUNT(*) FROM pages",
            "UNION ALL SELECT 'users', COUNT(*) FROM users",
            "UNION ALL SELECT 'comments', COUNT(*) FROM comments",
            "UNION ALL SELECT 'replies', COUNT(*) FROM comments WHERE rid != 0;"]
    (out_dir / "99-post.sql").write_text("\n".join(tail) + "\n", encoding="utf-8", newline="\n")
    files.append("99-post.sql")

    # 顺序执行脚本
    # -y 必须加：remote 执行会先问一次确认，非交互环境下不加会直接报错退出
    # 每个文件重试 3 次：跨境的 fetch 偶发失败（实测在 03-comments-007 断过一次）
    lines = [
        "#!/usr/bin/env bash",
        "set -uo pipefail",
        f"DB={args.db}",
        "MODE=${1:---remote}",
        "export WRANGLER_SEND_METRICS=false",
        "WR=\"node node_modules/wrangler/bin/wrangler.js\"",
        '[ -f "node_modules/wrangler/bin/wrangler.js" ] || WR="npx wrangler"',
        "",
        "FAILED=0",
    ]
    for f in files:
        lines.append(f'echo "==> {f}"')
        lines.append("ok=0")
        lines.append("for try in 1 2 3; do")
        lines.append(f'  if $WR d1 execute "$DB" "$MODE" -y --file={out_dir.as_posix()}/{f} >/dev/null 2>&1; then ok=1; break; fi')
        lines.append('  echo "    第 $try 次失败，重试（跨境 fetch 偶发失败）"')
        lines.append("  sleep 3")
        lines.append("done")
        lines.append('if [ "$ok" = "1" ]; then echo "    ok"; else echo "    FAIL"; FAILED=1; fi')
    lines.append("")
    lines.append('echo "核对条数："')
    lines.append(
        f'$WR d1 execute "$DB" "$MODE" -y --command '
        '"SELECT (SELECT COUNT(*) FROM sites) AS sites, (SELECT COUNT(*) FROM pages) AS pages,'
        ' (SELECT COUNT(*) FROM users) AS users, (SELECT COUNT(*) FROM comments) AS comments,'
        ' (SELECT COUNT(*) FROM comments WHERE rid!=0) AS replies;"'
    )
    lines.append("")
    lines.append('[ "$FAILED" = "0" ] && echo "全部导入成功" || echo "有文件失败，看上面 FAIL 那几行"')
    (out_dir / "run.sh").write_text("\n".join(lines) + "\n", encoding="utf-8", newline="\n")

    print(f"\n生成 {len(files)} 个 SQL 于 {out_dir}")
    print(f"顺序执行： bash {out_dir.as_posix()}/run.sh --remote")
    return 0


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