#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
merge-engine-data.py — 跨平台(Windows / Linux / macOS)把分散在多个 opencode 库里的
会话「合并」到引擎实际读写的那个库:既保留隔离后新建的会话,又找回历史会话。

为什么用 Python:只依赖 Python 自带的 sqlite3 模块,Windows 无需额外装 sqlite3.exe,
三个平台同一份脚本即可运行。

核心:
  - 自动确定「活库」= 引擎实际读写的库。判据:在所有「隔离引擎库」(路径里含
    openwork-engine-data、且不在 backup-* 里)中,取「含最新会话」的那个 —— 这必然是
    引擎正在写的库。这样自动绕开「productName 还是 appId 命名」「旧脚本留下的诱饵库」
    等差异;若一个隔离库都没有(非隔离版),退而取全局含最新会话的现存库。
  - 来源 = 找到的其它所有 opencode*.db(历史库 + 之前迁移/合并生成的 backup-*/ 里的库)。
  - 只增不删:按主键去重(INSERT OR IGNORE),冲突保留目标库(活库)的行 →
    隔离后新建的会话不会被旧历史覆盖。按两库公共列复制,兼容不同 opencode 版本的列差异。
  - 合并前自动备份目标库,可回滚。

用法:
  python3 merge-engine-data.py                 # 一步:自动选活库 + 并入其它所有库
  python3 merge-engine-data.py --dry-run        # 只打印计划
  python3 merge-engine-data.py --yes
  python3 merge-engine-data.py --into <db> --from <db> [--from <db> ...]   # 手动指定
  (Windows 用 py merge-engine-data.py ... ;若 python3 不可用改用 python)
"""
import argparse
import os
import shutil
import sqlite3
import subprocess
import sys
import time
from pathlib import Path

PRODUCT_NAME = "DeepWorks"

# 仅合并“会话/项目内容”表,按外键依赖顺序;绝不碰迁移账本 / 账号登录态 / 事件 / 工作区注册表。
MERGE_TABLES = [
    "project", "session", "message", "part",
    "session_message", "todo", "session_share", "permission",
]


def log(msg=""):
    print(msg, flush=True)


def candidate_roots():
    home = Path.home()
    roots = []
    if sys.platform == "win32":
        v = os.environ.get("APPDATA")            # Electron userData on Windows = %APPDATA%\<name>
        if v:
            roots.append(Path(v))
        roots += [home / ".local" / "share", home / ".config"]
    elif sys.platform == "darwin":
        roots += [home / "Library" / "Application Support", home / ".local" / "share", home / ".config"]
    else:  # linux / other
        roots += [home / ".local" / "share", home / ".config"]
    for env in ("XDG_DATA_HOME", "XDG_CONFIG_HOME"):
        v = os.environ.get(env)
        if v:
            roots.append(Path(v))
    eng = os.environ.get("OPENWORK_ENGINE_DATA_DIR")
    if eng:
        roots.append(Path(eng))
    # 去重并仅保留存在的目录
    seen, out = set(), []
    for r in roots:
        rp = str(r)
        if rp in seen:
            continue
        seen.add(rp)
        if r.is_dir():
            out.append(r)
    return out


# 引擎库的目录结构是固定的(runtime.mjs:<root>/openwork-engine-data/xdg/data/opencode),
# 因此用精确 glob 直接命中,避免在 %APPDATA% 这类大目录里做深度遍历(Windows 上更快更稳)。
_GLOB_PATTERNS = [
    "xdg/data/opencode/opencode*.db",                              # OPENWORK_ENGINE_DATA_DIR 直指根
    "xdg/data/opencode/backup-*/opencode*.db",
    "*/openwork-*-data/xdg/data/opencode/opencode*.db",            # <name>/openwork-engine-data(或 -dev-data)
    "*/openwork-*-data/xdg/data/opencode/backup-*/opencode*.db",
    "opencode/opencode*.db",                                       # 共享 opencode 数据目录
    "opencode/backup-*/opencode*.db",
]


def find_dbs(roots):
    found = set()
    for root in roots:
        for pat in _GLOB_PATTERNS:
            try:
                for p in root.glob(pat):
                    if p.is_file():
                        found.add(str(p))
            except Exception:
                pass
    return sorted(found)


def _scalar(path, sql):
    try:
        con = sqlite3.connect(path)
        try:
            cur = con.execute(sql)
            row = cur.fetchone()
            return row[0] if row and row[0] is not None else 0
        finally:
            con.close()
    except Exception:
        return 0


def sessions_of(path):
    return _scalar(path, "select count(*) from session")


def messages_of(path):
    return _scalar(path, "select count(*) from message")


def newest_of(path):
    return _scalar(path, "select coalesce(max(time_updated),0) from session")


def is_backup(path):
    return any(part.startswith("backup-") for part in Path(path).parts)


def is_engine_iso(path):
    p = Path(path)
    return "openwork-engine-data" in p.parts and p.name.startswith("opencode") and p.suffix == ".db"


def table_columns(con, schema, table):
    cur = con.execute('PRAGMA "%s".table_info("%s")' % (schema, table))
    return [r[1] for r in cur.fetchall()]


def table_exists(con, schema, table):
    cur = con.execute(
        'select 1 from "%s".sqlite_master where type=\'table\' and name=? limit 1' % schema,
        (table,),
    )
    return cur.fetchone() is not None


def app_running():
    try:
        if sys.platform == "win32":
            out = subprocess.run(["tasklist"], capture_output=True, text=True, timeout=10).stdout
            return PRODUCT_NAME.lower() in out.lower()
        else:
            r = subprocess.run(["pgrep", "-f", PRODUCT_NAME], capture_output=True, timeout=10)
            return r.returncode == 0
    except Exception:
        return False  # 检测不了就不挡,交给用户判断


def merge_one(target_db, src_db):
    """把 src_db 的会话内容并入 target_db(只增不删)。返回 True/False。"""
    con = sqlite3.connect(target_db, isolation_level=None)  # autocommit,手动控制事务
    try:
        con.execute("PRAGMA foreign_keys=OFF")
        con.execute("ATTACH DATABASE ? AS src", (src_db,))
        con.execute("BEGIN")
        try:
            for t in MERGE_TABLES:
                if not table_exists(con, "main", t) or not table_exists(con, "src", t):
                    continue
                main_cols = table_columns(con, "main", t)
                src_cols = set(table_columns(con, "src", t))
                cols = [c for c in main_cols if c in src_cols]
                if not cols:
                    continue
                collist = ",".join('"%s"' % c for c in cols)
                con.execute(
                    'INSERT OR IGNORE INTO main."%s" (%s) SELECT %s FROM src."%s"'
                    % (t, collist, collist, t)
                )
            con.execute("COMMIT")
        except Exception:
            con.execute("ROLLBACK")
            raise
        finally:
            con.execute("DETACH DATABASE src")
        return True
    finally:
        con.close()


def main():
    ap = argparse.ArgumentParser(add_help=True, description="合并 opencode 会话库到引擎活库")
    ap.add_argument("--dry-run", action="store_true", help="只打印计划,不改动")
    ap.add_argument("--yes", "-y", action="store_true", help="跳过确认")
    ap.add_argument("--into", help="手动指定目标活库(opencode.db)")
    ap.add_argument("--from", dest="froms", action="append", default=[], help="手动指定来源库(可多次)")
    args = ap.parse_args()

    all_dbs = find_dbs(candidate_roots())

    # --- 确定目标活库 ---
    into = args.into
    reason = ""
    if not into:
        iso_best, iso_ts, any_best, any_ts = None, -1, None, -1
        for db in all_dbs:
            if is_backup(db):
                continue
            ts = newest_of(db)
            if ts > any_ts:
                any_ts, any_best = ts, db
            if is_engine_iso(db) and ts > iso_ts:
                iso_ts, iso_best = ts, db
        if iso_best:
            into, reason = iso_best, "隔离引擎库中含最新会话者(= 引擎实际读写的库;已跳过诱饵库/备份)"
        elif any_best:
            into, reason = any_best, "含全局最新会话的现存库(非隔离版/特殊情况)"

    # --- 来源 = 其它所有库 ---
    froms = list(args.froms) if args.froms else [db for db in all_dbs if db != into]

    log("==> 合并计划")
    if not into or not Path(into).is_file():
        log("找不到可作为目标的现存 opencode 库(只发现备份或一个都没有)。")
        log("请先打开一次 %s 再完全退出,让引擎用正确 schema 建好隔离库,然后重跑;" % PRODUCT_NAME)
        log("或用 --into <路径>/opencode.db 手动指定引擎实际读写的库。")
        sys.exit(1)
    log("  目标库(活库): %s" % into)
    if reason:
        log("  选择依据:     %s" % reason)
    log("  目标库现有:   %s 会话 / %s 消息" % (sessions_of(into), messages_of(into)))
    froms = [f for f in froms if Path(f).is_file() and os.path.abspath(f) != os.path.abspath(into)]
    if not froms:
        log("\n只发现目标库这一个库,没有其它可并入的来源。无需合并。")
        sys.exit(0)
    log("  将并入以下来源库(只增不删,按主键去重,冲突保留目标库):")
    for db in froms:
        log("    + %s   [%s 会话 / %s 消息]" % (db, sessions_of(db), messages_of(db)))

    if args.dry_run:
        log("\n(--dry-run) 未做任何改动。")
        return

    if os.environ.get("OPENWORK_MIGRATE_SKIP_RUNNING_CHECK") != "1" and app_running():
        log("\n检测到 %s 仍在运行。请完全退出 App 后再运行本脚本。" % PRODUCT_NAME)
        sys.exit(1)

    if not args.yes:
        try:
            reply = input("\n确认合并？目标库会先备份。[y/N] ").strip().lower()
        except EOFError:
            reply = ""
        if reply not in ("y", "yes"):
            log("已取消。")
            return

    # --- 备份目标库 ---
    ts = time.strftime("%Y%m%d-%H%M%S")
    backup = Path(into).parent / ("backup-merge-" + ts)
    backup.mkdir(parents=True, exist_ok=True)
    log("\n==> 备份目标库 → %s" % backup)
    for suffix in ("", "-wal", "-shm"):
        f = Path(str(into) + suffix)
        if f.is_file():
            shutil.copy2(str(f), str(backup / f.name))
            log("  备份 %s" % f.name)

    # --- 逐个合并 ---
    for src in froms:
        log("\n==> 合并 %s" % src)
        try:
            merge_one(into, src)
            log("  完成 → 目标库现有 %s 会话 / %s 消息" % (sessions_of(into), messages_of(into)))
        except Exception as e:
            log("  合并失败(已回滚该来源):%s" % e)
            log("  目标库未损坏,可用 %s 还原。" % backup)

    # --- 外键自检 ---
    log("\n==> 外键自检(应无输出表示正常)")
    try:
        con = sqlite3.connect(into)
        rows = con.execute("PRAGMA foreign_key_check").fetchall()
        con.close()
        if rows:
            for r in rows[:5]:
                log("  ⚠ %s" % (r,))
            log("  存在外键悬挂(通常是来源库缺对应 project/session 行);引擎一般仍可读,如异常请用备份还原。")
        else:
            log("  无外键问题。")
    except Exception as e:
        log("  自检跳过:%s" % e)

    log("\n==> 完成")
    log("  已合并到活库:%s(现 %s 会话 / %s 消息)。重开 %s 查看。"
        % (into, sessions_of(into), messages_of(into), PRODUCT_NAME))
    log("\n回滚(还原合并前的目标库):把 %s 里的 opencode.db* 拷回 %s"
        % (backup, Path(into).parent))


if __name__ == "__main__":
    main()