zcode_import.py

#!/usr/bin/env python3
# ---------------------------------------------------------------------------
# Reference implementation: ZCode's official `restore-legacy-sessions` plugin
#   @zcode/restore-legacy-sessions-plugin 0.1.0  (c) Z.ai
#   MIT licensed, per package.json / .zcode-plugin/plugin.json.
#   https://opensource.org/license/MIT
# This tool drives those official scripts directly when available (preferred)
# and otherwise falls back to a built-in Python port of `restore-conversation.mjs`.
# ---------------------------------------------------------------------------
"""Import ZCode restore-legacy-sessions snapshots into ZCode's stores.

Usage:
  python zcode_import.py list [--legacy-dir DIR] [--agent A] [--workspace W]
                              [--query Q] [--conversation ID] [--limit N]
                              [--json] [--task-index PATH] [--cli-db PATH]
  python zcode_import.py status [same filter/legacy-dir/db options]
  python zcode_import.py import <snapshot path, ...> | all
                              [--legacy-dir DIR] [--task-index PATH]
                              [--cli-db PATH] [--dry-run] [--show] [--no-show]
                              [--backup]
  python zcode_import.py unarchive <taskId, ...> | all [--task-index PATH]
--show:  un-hide every restored task in the task index (archived=0, deleted=0,
         pinned=0, unread_at=now). ZCode's task list only shows deleted=0 rows
         (and the grouped view needs archived=0), and its auto-archive sweeps
         old completed rows, so freshly restored legacy sessions can vanish
         from the app. This is ON BY DEFAULT; --no-show turns it off.
target:  a snapshot .json path, an ordinal (see # below), a directory (workspace
         dir or export root) whose sess_*.json snapshots are imported, or 'all'
# (list) and import <n>: the '#' column of `list` is a GLOBAL, filter-independent
         ordinal computed over the full legacy-dir snapshot set (updatedAt
         descending), the same numbering `import <n>` uses. It never depends on
         --agent/--workspace/--query/--limit, so filtered list views may show
         non-contiguous numbers (intended). `import <n>` restores that snapshot
         no matter what filters were active.
unarchive: un-hide task rows already present in the task index (no import).
--backup: retain/generate a pre-write timestamped backup copy of both DBs.
         OFF BY DEFAULT. With the official backend the scripts always make a
         backup on import; without --backup our tool deletes the copies this
         import produced when it finishes (as if there had been no backup), on
         success only. With the python backend no backup is made at all unless
         --backup is given.

Backend selection (hybrid):
  * official  (default when available): shell out to the installed plugin
    scripts `restore-legacy-sessions/.../scripts/{scan-legacy-sessions.mjs,
    restore-conversation.mjs}` via `node` (>= 22.5). The plugin's LLM-driven
    "pick a session" skill is NOT used; the scripts are invoked directly.
  * python    (fallback): built-in port of `restore-conversation.mjs`,
    no Node required. Force either choice with $ZCODE_IMPORT_BACKEND
    (official | python); "official" falls back when tools are missing.

Where the official scripts are looked up:
  1. $ZCODE_LEGACY_SESSIONS_SCRIPTS          (dir containing both .mjs files)
  2. ~/.zcode/cli/plugins/cache/zcode-plugins-official/restore-legacy-sessions/<v>/skills/restore-legacy-sessions/scripts
  3. ProgramFiles\\ZCode\\resources\\glm\\packages\\restore-legacy-sessions-plugin\\skills\\restore-legacy-sessions\\scripts

Defaults:
  --legacy-dir  ~/.zcode/v2/sessions           ($ZCODE_LEGACY_DIR)
  --task-index  ~/.zcode/v2/tasks-index.sqlite ($ZCODE_TASK_INDEX)
  --cli-db      ~/.zcode/cli/db/db.sqlite      ($ZCODE_DB)
  scripts dir   see above                      ($ZCODE_LEGACY_SESSIONS_SCRIPTS)
  backend       official when available        ($ZCODE_IMPORT_BACKEND)
Requires: python 3.7+ (stdlib only); node >= 22.5 only for the official backend.
"""

from __future__ import annotations

import glob
import json
import os
import re
import shutil
import sqlite3
import subprocess
import sys
import time

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

HOME = os.path.expanduser("~")
LEGACY_DIR = (
    os.environ.get("ZCODE_LEGACY_DIR")
    or os.path.join(HOME, ".zcode", "v2", "sessions")
)
TASK_INDEX = (
    os.environ.get("ZCODE_TASK_INDEX")
    or os.path.join(HOME, ".zcode", "v2", "tasks-index.sqlite")
)
CLI_DB = (
    os.environ.get("ZCODE_DB")
    or os.path.join(HOME, ".zcode", "cli", "db", "db.sqlite")
)

ZCODE_PROVIDER = "glm"


def ts(ms):
    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ms / 1000))


def json_object(value):
    return value if isinstance(value, dict) else {}


def json_text(value):
    if value is None:
        return ""
    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False,
                                                           separators=(",", ":"))


def jdumps(value):
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def optional_string(value):
    return value if isinstance(value, str) and value.strip() else None


def optional_number(value):
    return value if isinstance(value, (int, float)) and value == value else None


def remove_undefined(obj):
    return {k: v for k, v in (obj or {}).items() if v is not None}


def backup_db(path):
    stamp = (
        time.strftime("%Y-%m-%dT%H-%M-%S", time.gmtime())
        + f".{int(time.time() * 1000) % 1000:03d}Z"
    )
    backup = f"{path}.bak-{stamp}"
    shutil.copyfile(path, backup)
    return backup


_BACKUP_RE = re.compile(r"\.bak-\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}[-.]\d{3}Z$")


def list_dir_backup_files(db_path):
    """Absolute paths of backup copies of a db file currently on disk.

    Matches only the main `<db>.bak-<timestamp>` copy created by either the
    official script or our python backend. Sidecar `-shm`/`-wal` files and any
    unrelated `.bak-*` names are deliberately excluded.
    """
    d = os.path.dirname(db_path) or "."
    base = os.path.basename(db_path)
    out = []
    for name in os.listdir(d):
        if not name.startswith(base + ".bak"):
            continue
        if not _BACKUP_RE.search(name):
            continue
        out.append(os.path.join(d, name))
    return out


def cleanup_official_backups(before, task_index, cli_db):
    """Delete the official backups created by this import run.

    The official restore script always makes timestamped copies of both DBs.
    Without --backup we remove exactly the copies that appeared during this
    import (before/after set difference), leaving history untouched. Only the
    main copy is deleted; -shm/-wal sidecars and live DBs are never touched.
    """
    after = set(list_dir_backup_files(task_index))
    after |= set(list_dir_backup_files(cli_db))
    new = sorted(after - set(before))
    for path in new:
        try:
            os.remove(path)
            print(f"Removed official backup: {path}")
        except OSError as exc:
            print(f"Could not remove {path}: {exc}", file=sys.stderr)
    if not new:
        print("No official backup retained (no new backup created).")
def ensure_existing_db(path, label):
    if not os.path.exists(path):
        raise SystemExit(f"{label} not found: {path}")
    if not os.path.isfile(path) or os.path.getsize(path) == 0:
        raise SystemExit(f"{label} is not a populated sqlite file: {path}")


# ---------------------------------------------------------------- snapshot scan

def collect_json_files(root):
    if not os.path.isdir(root):
        return []
    files = []
    for ws_name in sorted(os.listdir(root)):
        ws_dir = os.path.join(root, ws_name)
        if not os.path.isdir(ws_dir):
            continue
        for fn in sorted(os.listdir(ws_dir)):
            if fn.endswith(".json") and not fn.endswith(".deleted.json"):
                files.append(os.path.join(ws_dir, fn))
    return files


def collect_dir_snapshots(root):
    """Find ZCode snapshot files under a directory (recursive).

    A target directory may be a single workspace dir (flat <sess>.json) or an
    export root (<hash>/<sess>.json). Only snapshot-shaped files are picked
    (basename starts with sess_ and ends with .json, skipping *.deleted.json)
    so stray md/json/opencode exports are ignored. This is CLI-level
    convenience: the resolved file list is fed identically to both backends.
    """
    files = [p for p in glob.glob(os.path.join(root, "**", "*.json"),
                                  recursive=True)
             if os.path.isfile(p)]
    out = []
    for p in files:
        name = os.path.basename(p)
        if name.startswith("sess_") and not name.endswith(".deleted.json"):
            out.append(p)
    return sorted(out)


def build_global_ordinals(legacy_dir):
    """Fixed, filter-independent ordinal for every snapshot under legacy_dir.

    Ordinals are computed once on the full two-level candidate set (same
    collect_json_files that the official scanner iterates) sorted by updatedAt
    descending. This is the single authoritative numbering used both by the
    '#' column of `list` and by `import <n>`, so a number never depends on the
    filters/limit active in some particular list view (rows may therefore show
    non-contiguous '#' when filtered, which is intended).
    """
    candidates = [c for p in collect_json_files(legacy_dir)
                  for c in [read_snapshot(p)] if not c.get("error")]
    candidates.sort(key=lambda c: -(c.get("updatedAt") or 0))
    return {os.path.abspath(c["filePath"]): i for i, c in enumerate(candidates, 1)}


def first_user_content(messages):
    for m in messages:
        if isinstance(m, dict) and m.get("role") == "user":
            c = m.get("content")
            return c if isinstance(c, str) else None
    return None


def summarize_text(value):
    if not isinstance(value, str):
        return ""
    s = re.sub(r"\s+", " ", value).strip()
    return s[:80] + ("..." if len(s) > 80 else "")


def finite_number(value):
    return value if isinstance(value, (int, float)) else None


def first_timestamp(messages):
    for m in messages:
        if isinstance(m, dict):
            t = finite_number(m.get("timestamp"))
            if t is not None:
                return int(t)
    return None


def last_timestamp(messages):
    for m in reversed(messages):
        if isinstance(m, dict):
            t = finite_number(m.get("timestamp"))
            if t is not None:
                return int(t)
    return None


def read_snapshot(file_path):
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            raw = json.load(f)
        messages = raw.get("messages") if isinstance(raw.get("messages"), list) else []
        meta = json_object(raw.get("meta"))
        legacy_task_id = str(meta.get("taskId") or os.path.splitext(
            os.path.basename(file_path))[0])
        acp_session_id = meta.get("acpSessionId") if isinstance(
            meta.get("acpSessionId"), str) else ""
        restored_task_id = acp_session_id or legacy_task_id
        workspace_path = meta.get("workspacePath") if isinstance(
            meta.get("workspacePath"), str) else ""
        provider = meta.get("provider") if isinstance(meta.get("provider"), str) else "unknown"
        title = meta.get("title") if isinstance(meta.get("title"), str) and meta.get("title").strip() \
            else summarize_text(first_user_content(messages)) or "Untitled session"
        created_at = finite_number(meta.get("createdAt"))
        if created_at is None:
            created_at = first_timestamp(messages)
        if created_at is None:
            created_at = 0
        updated_at = finite_number(meta.get("updatedAt"))
        if updated_at is None:
            updated_at = last_timestamp(messages)
        if updated_at is None:
            updated_at = created_at
        visible_text = "\n".join(
            m.get("content") for m in messages
            if isinstance(m.get("content"), str))
        result = {
            "filePath": file_path,
            "workspaceHash": os.path.basename(os.path.dirname(file_path)),
            "legacyTaskId": legacy_task_id,
            "acpSessionId": acp_session_id,
            "restoredTaskId": restored_task_id,
            "workspacePath": workspace_path,
            "provider": provider,
            "model": meta.get("model") if isinstance(meta.get("model"), str) else None,
            "title": title,
            "createdAt": int(created_at),
            "updatedAt": int(updated_at),
            "messageCount": len(messages),
            "userMessageCount": sum(1 for m in messages if m.get("role") == "user"),
            "assistantMessageCount": sum(1 for m in messages if m.get("role") == "assistant"),
            "visibleText": visible_text,
        }
        wi = meta.get("workspaceIdentity")
        if isinstance(wi, str):
            result["workspaceIdentity"] = wi
        return result
    except Exception as exc:
        return {"filePath": file_path, "error": str(exc)}


def open_readonly(path):
    if not os.path.exists(path):
        return None
    if os.path.getsize(path) == 0:
        return None
    try:
        return sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    except Exception:
        return None


def count(db, sql, args):
    if db is None:
        return 0
    try:
        return db.execute(sql, args).fetchone()[0]
    except Exception:
        return 0


def attach_store_status(candidates, task_index, cli_db):
    ti = open_readonly(task_index)
    cd = open_readonly(cli_db)
    task_restored = (
        "SELECT COUNT(*) FROM tasks WHERE workspace_path=? AND task_id=?",
        lambda c: (c["workspacePath"], c["restoredTaskId"]),
    )
    task_legacy = (
        "SELECT COUNT(*) FROM tasks WHERE workspace_path=? AND task_id=?",
        lambda c: (c["workspacePath"], c["legacyTaskId"]),
    )
    cli_session = (
        "SELECT COUNT(*) FROM session WHERE id=?",
        lambda c: (c["restoredTaskId"],),
    )
    for c in candidates:
        if c.get("error"):
            continue
        task_index_restored = count(ti, task_restored[0], task_restored[1](c))
        task_index_legacy = count(ti, task_legacy[0], task_legacy[1](c))
        cli_session_count = count(cd, cli_session[0], cli_session[1](c))
        c["store"] = {
            "cliDb": "present" if cli_session_count > 0 else "missing",
            "taskIndex": "present" if task_index_restored > 0
                else ("legacy-task-id" if task_index_legacy > 0 else "missing"),
        }
        c["restoreState"] = restore_state(c["store"])
    if ti is not None:
        ti.close()
    if cd is not None:
        cd.close()


def restore_state(store):
    cli_present = store["cliDb"] == "present"
    task_present = store["taskIndex"] == "present"
    if cli_present and task_present:
        return "ready"
    if not cli_present and task_present:
        return "needs-cli-db"
    if cli_present and not task_present:
        return "needs-task-index"
    return "needs-full-import"


def filter_candidates(candidates, args):
    query = (args.get("query") or "").lower()
    conv = (args.get("conversation") or "").lower()
    out = []
    for c in candidates:
        if c.get("error"):
            continue
        if args.get("agent") and args["agent"] != "all" and c["provider"] != args["agent"]:
            continue
        if args.get("workspace") and c["workspacePath"] != args["workspace"]:
            continue
        if conv and conv not in {
            v.lower() for v in (
                (c["restoredTaskId"] or ""),
                (c["legacyTaskId"] or ""),
                (c["acpSessionId"] or "")) if v
        }:
            continue
        if query:
            haystack = " ".join(filter(None, [
                c["provider"], c["workspacePath"], c["title"],
                c["restoredTaskId"], c["legacyTaskId"], c["acpSessionId"],
                c["visibleText"]])).lower()
            if query not in haystack:
                continue
        out.append(c)
    return out


def group_by(items, key_fn):
    groups = {}
    for item in items:
        key = key_fn(item)
        groups.setdefault(key, []).append(item)
    return groups


def summarize_groups(candidates, key_fn):
    rows = []
    for key, items in group_by(candidates, key_fn).items():
        states = [i["restoreState"] for i in items if not i.get("error")]
        rows.append({
            "key": key,
            "count": len(items),
            "workspaces": len({i.get("workspacePath") for i in items}),
            "ready": states.count("ready"),
            "needsCliDb": states.count("needs-cli-db"),
            "needsTaskIndex": states.count("needs-task-index"),
            "needsFullImport": states.count("needs-full-import"),
            "latestUpdatedAt": max((i.get("updatedAt") or 0 for i in items), default=0),
        })
    return sorted(rows, key=lambda r: (-r["count"], r["key"]))


def md_table(headers, rows):
    esc = lambda v: "" if v is None else str(v).replace("|", "\\|").replace("\n", " ")
    lines = ["| " + " | ".join(headers) + " |",
             "| " + " | ".join(["---"] * len(headers)) + " |"]
    lines += ["| " + " | ".join(esc(v) for v in row) + " |" for row in rows]
    return "\n".join(lines)


# -------------------------------------------------------------- restore side

def message_content(message):
    c = message.get("content")
    if isinstance(c, str):
        return c
    if c is None:
        return ""
    return str(c)


def message_completed_at(message, message_timestamp):
    t = finite_number(message.get("timestamp"))
    if t is not None:
        return int(t) + int(finite_number(message.get("durationMs")) or 0)
    return message_timestamp


def legacy_tool_name(tool):
    raw = json_object(tool.get("raw"))
    meta = json_object(raw.get("_meta"))
    claude = json_object(meta.get("claudeCode"))
    return (optional_string(claude.get("toolName"))
            or optional_string(tool.get("toolName"))
            or optional_string(tool.get("name"))
            or optional_string(tool.get("kind"))
            or optional_string(tool.get("title"))
            or "legacy_tool")


def legacy_tool_title(tool, tool_name):
    return optional_string(tool.get("title")) or tool_name


def legacy_tool_call_id(message_id, tool_index, tool):
    return (optional_string(tool.get("id"))
            or optional_string(tool.get("callID"))
            or optional_string(tool.get("callId"))
            or f"call_legacy_{message_id}_{tool_index}")


def legacy_tool_metadata(tool):
    return remove_undefined({
        "legacyRestore": True,
        "legacyKind": optional_string(tool.get("kind")),
        "legacyStatus": optional_string(tool.get("status")),
    })


def build_tool_part_data(tool, message_id, tool_index, message_timestamp, completed_at):
    tool_name = legacy_tool_name(tool)
    state = {"input": json_object(tool.get("input"))}
    status = optional_string(tool.get("status")) or "completed"
    started_at = optional_number(tool.get("startedAt"))
    if started_at is None:
        started_at = message_timestamp
    ended_at = optional_number(tool.get("completedAt"))
    if ended_at is None:
        ended_at = completed_at
    metadata = legacy_tool_metadata(tool)
    if status == "running":
        state.update({
            "status": "running",
            "title": legacy_tool_title(tool, tool_name),
            "metadata": metadata,
            "time": {"start": started_at},
        })
    elif status == "pending":
        state.update({
            "status": "pending",
            "raw": json_text(tool.get("raw") or tool),
        })
    elif status in ("failed", "error", "denied"):
        state.update({
            "status": "error",
            "error": json_text(tool.get("error") or tool.get("output") or status),
            "metadata": metadata,
            "time": {"start": started_at, "end": ended_at},
        })
    else:
        state.update({
            "status": "completed",
            "output": json_text(tool.get("output")),
            "title": legacy_tool_title(tool, tool_name),
            "metadata": metadata,
            "time": {"start": started_at, "end": ended_at},
        })
    return {
        "type": "tool",
        "callID": legacy_tool_call_id(message_id, tool_index, tool),
        "tool": tool_name,
        "state": state,
    }


def build_text_part_data(text, message_timestamp, completed_at):
    return {"type": "text", "text": text,
            "time": {"start": message_timestamp, "end": completed_at}}


def build_reasoning_part_data(text, message_timestamp, completed_at):
    return {"type": "reasoning", "text": text,
            "time": {"start": message_timestamp, "end": completed_at}}


def build_part_data_list(message, message_id, message_timestamp):
    content = message_content(message)
    completed_at = message_completed_at(message, message_timestamp)
    tools = message.get("tools") if isinstance(message.get("tools"), list) else []
    persisted = message.get("parts") if isinstance(message.get("parts"), list) else []
    part_data = []

    if message.get("role") == "assistant" and persisted:
        for part in persisted:
            if part.get("type") == "content" and isinstance(part.get("content"), str) \
                    and part["content"]:
                part_data.append(build_text_part_data(
                    part["content"], message_timestamp, completed_at))
            elif part.get("type") == "thought" and isinstance(part.get("content"), str) \
                    and part["content"]:
                part_data.append(build_reasoning_part_data(
                    part["content"], message_timestamp, completed_at))
            elif part.get("type") == "tool-call" and isinstance(part.get("toolIndex"), int):
                tool_index = part["toolIndex"]
                if 0 <= tool_index < len(tools):
                    part_data.append(build_tool_part_data(
                        tools[tool_index], message_id, tool_index,
                        message_timestamp, completed_at))
    else:
        thought = optional_string(message.get("thought"))
        if message.get("role") == "assistant" and thought:
            part_data.append(build_reasoning_part_data(
                thought, message_timestamp, completed_at))
        if content:
            part_data.append(build_text_part_data(content, message_timestamp, completed_at))
        for tool_index, tool in enumerate(tools):
            part_data.append(build_tool_part_data(
                tool, message_id, tool_index, message_timestamp, completed_at))

    has_visible_text = any(
        (p["type"] in ("text", "reasoning") and p.get("text")) for p in part_data)
    if not has_visible_text and content:
        part_data.append(build_text_part_data(content, message_timestamp, completed_at))
    return part_data


def project_id_from_path(workspace_path):
    return "proj_" + re.sub(r"^/", "", workspace_path).replace("/", "-")


def workspace_key(workspace_path, workspace_identity):
    return (workspace_identity or "").strip() or workspace_path


def build_restore_plan(snapshot_path):
    with open(snapshot_path, "r", encoding="utf-8") as f:
        raw = json.load(f)
    meta = json_object(raw.get("meta"))
    messages = raw.get("messages") if isinstance(raw.get("messages"), list) else []

    legacy_task_id = str(meta.get("taskId") or os.path.splitext(
        os.path.basename(snapshot_path))[0])
    acp_session_id = optional_string(meta.get("acpSessionId")) or ""
    restored_task_id = acp_session_id or legacy_task_id
    workspace_path = optional_string(meta.get("workspacePath"))
    if not workspace_path:
        raise SystemExit(f"Snapshot has no workspacePath: {snapshot_path}")
    workspace_identity = optional_string(meta.get("workspaceIdentity")) or ""
    workspace_hash = os.path.basename(os.path.dirname(snapshot_path))
    source_provider = optional_string(meta.get("provider")) or "unknown"
    title = optional_string(meta.get("title")) or "Untitled session"
    mode = optional_string(meta.get("mode")) or "build"
    model = optional_string(meta.get("model")) or ""
    trace_id = optional_string(meta.get("traceId")) or f"zcode-{restored_task_id}"
    created_at = finite_number(meta.get("createdAt"))
    created_at = int(created_at) if created_at is not None else int(time.time() * 1000)
    updated_at = finite_number(meta.get("updatedAt"))
    updated_at = int(updated_at) if updated_at is not None else created_at

    normalized_meta = remove_undefined({
        **meta,
        "taskId": restored_task_id,
        "traceId": trace_id,
        "title": title,
        "workspacePath": workspace_path,
        "workspaceIdentity": workspace_identity or None,
        "createdAt": created_at,
        "updatedAt": updated_at,
        "mode": mode,
        "model": model or None,
        "thoughtLevel": optional_string(meta.get("thoughtLevel")),
        "runtimeEpoch": finite_number(meta.get("runtimeEpoch")),
        "provider": ZCODE_PROVIDER,
        "legacyTaskId": legacy_task_id,
        "restoredTaskId": restored_task_id,
        "workspaceHash": workspace_hash,
        "messageCount": len(messages),
        "migrationSource": None,
    })

    return {
        "snapshotPath": snapshot_path,
        "raw": raw,
        "meta": meta,
        "messages": messages,
        "legacyTaskId": legacy_task_id,
        "restoredTaskId": restored_task_id,
        "workspacePath": workspace_path,
        "workspaceIdentity": workspace_identity,
        "workspaceHash": workspace_hash,
        "workspaceKey": workspace_key(workspace_path, workspace_identity),
        "projectId": project_id_from_path(workspace_path),
        "sourceProvider": source_provider,
        "title": title,
        "mode": mode,
        "model": model,
        "traceId": trace_id,
        "createdAt": created_at,
        "updatedAt": updated_at,
        "searchableText": "\n".join(m.get("content") for m in messages
                                    if isinstance(m.get("content"), str)),
        "metaJson": jdumps(normalized_meta),
    }


def print_plan(plan, dry_run=False):
    print("=== Importing snapshot ===")
    print(f"  Title:          {plan['title']}")
    print(f"  SourceProvider: {plan['sourceProvider']}")
    print(f"  RuntimeProvider: {ZCODE_PROVIDER}")
    print(f"  Model:          {plan['model']}")
    print(f"  Workspace:      {plan['workspacePath']}")
    print(f"  Messages:       {len(plan['messages'])}")
    print(f"  restoredTaskId: {plan['restoredTaskId']}")
    print(f"  legacyTaskId:   {plan['legacyTaskId']}")
    print(f"  workspaceHash:  {plan['workspaceHash']}")
    print("  migrationSource: <empty>")
    print(f"  Dry run:        {dry_run}")
    print()


def write_task_index(db, plan):
    sql = """
    INSERT INTO tasks
      (workspace_key, workspace_path, workspace_identity, task_id, title,
       task_status, provider, mode, model, migration_source, forked_from_task_id,
       created_at, updated_at, unread_at, meta_json, searchable_text)
    VALUES
      (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?)
    ON CONFLICT(workspace_key, task_id) DO UPDATE SET
      workspace_path = excluded.workspace_path,
      workspace_identity = excluded.workspace_identity,
      title = CASE
        WHEN tasks.title_overridden = 1 THEN tasks.title
        ELSE excluded.title
      END,
      task_status = excluded.task_status,
      provider = excluded.provider,
      mode = excluded.mode,
      model = excluded.model,
      migration_source = NULL,
      forked_from_task_id = excluded.forked_from_task_id,
      created_at = MIN(tasks.created_at, excluded.created_at),
      updated_at = MAX(tasks.updated_at, excluded.updated_at),
      unread_at = tasks.unread_at,
      searchable_text = CASE
        WHEN tasks.updated_at > excluded.updated_at THEN tasks.searchable_text
        ELSE excluded.searchable_text
      END,
      meta_json = excluded.meta_json
    """
    forked = optional_string(plan["meta"].get("forkedFromTaskId"))
    status = optional_string(plan["meta"].get("status"))
    unread = finite_number(plan["meta"].get("unreadAt"))
    db.execute(sql, (
        plan["workspaceKey"],
        plan["workspacePath"],
        plan["workspaceIdentity"] or None,
        plan["restoredTaskId"],
        plan["title"],
        status,
        ZCODE_PROVIDER,
        plan["mode"],
        plan["model"] or None,
        forked,
        plan["createdAt"],
        plan["updatedAt"],
        int(unread) if unread is not None else None,
        plan["metaJson"],
        plan["searchableText"],
    ))


def write_session(db, plan):
    sql = """
    INSERT INTO session
      (id, project_id, workspace_id, slug, directory, path, title, version,
       time_created, time_updated, task_type, title_source, trace_id)
    VALUES
      (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    ON CONFLICT(id) DO UPDATE SET
      project_id = excluded.project_id,
      workspace_id = excluded.workspace_id,
      slug = excluded.slug,
      directory = excluded.directory,
      path = excluded.path,
      title = CASE
        WHEN session.title_source = 'custom' THEN session.title
        ELSE excluded.title
      END,
      version = excluded.version,
      time_created = MIN(session.time_created, excluded.time_created),
      time_updated = MAX(session.time_updated, excluded.time_updated),
      task_type = excluded.task_type,
      trace_id = COALESCE(session.trace_id, excluded.trace_id)
    """
    db.execute(sql, (
        plan["restoredTaskId"],
        plan["projectId"],
        plan["workspaceHash"],
        plan["restoredTaskId"],
        plan["workspacePath"],
        plan["workspacePath"],
        plan["title"],
        "0.14.5",
        plan["createdAt"],
        plan["updatedAt"],
        "interactive",
        "first_input",
        plan["traceId"],
    ))


def write_messages(db, plan):
    insert_message = """
    INSERT OR REPLACE INTO message
      (id, session_id, time_created, time_updated, data)
    VALUES (?, ?, ?, ?, ?)
    """
    insert_part = """
    INSERT OR REPLACE INTO part
      (id, message_id, session_id, time_created, time_updated, data)
    VALUES (?, ?, ?, ?, ?, ?)
    """
    db.execute(
        "DELETE FROM part WHERE session_id=? AND id LIKE ?",
        (plan["restoredTaskId"], f"part_legacy_{plan['restoredTaskId']}_%"))

    part_count = 0
    latest_user_message_id = ""
    for index, message in enumerate(plan["messages"]):
        message_timestamp = finite_number(message.get("timestamp"))
        message_timestamp = int(message_timestamp) if message_timestamp is not None \
            else plan["updatedAt"]
        completed_at = message_completed_at(message, message_timestamp)
        message_id = f"msg_legacy_{plan['restoredTaskId']}_{index}"
        role = "assistant" if message.get("role") == "assistant" else "user"

        base = {
            "role": role,
            "time": remove_undefined({
                "created": message_timestamp,
                "completed": completed_at if role == "assistant" else None,
            }),
            "agent": "zcode-agent",
        }
        if role == "assistant":
            message_data = {
                **base,
                "content": message_content(message),
                "modelID": message.get("model") or plan["model"] or None,
                "providerID": ZCODE_PROVIDER,
                "mode": plan["mode"],
                "path": {"cwd": plan["workspacePath"], "root": plan["workspacePath"]},
                "cost": 0,
                "tokens": {"input": 0, "output": 0, "reasoning": 0,
                           "cache": {"read": 0, "write": 0}},
                "finish": "stop",
            }
            if latest_user_message_id:
                message_data["parentID"] = latest_user_message_id
            tools = message.get("tools") if isinstance(message.get("tools"), list) else []
            if tools:
                message_data["toolCalls"] = [
                    {
                        "id": f"tool_{message_id}_{tool_index}",
                        "toolName": tool.get("title") or "",
                        "kind": tool.get("kind") or "other",
                        "status": tool.get("status") or "completed",
                        "input": tool.get("input") or {},
                        "output": tool.get("output") or "",
                        **({"raw": tool["raw"]} if tool.get("raw") else {}),
                    }
                    for tool_index, tool in enumerate(tools)
                ]
            for field in ("characterCount", "turnIndex"):
                if field in message and message[field] is not None:
                    message_data[field] = message[field]
        else:
            message_data = {
                **base,
                "content": message_content(message),
                "model": {
                    "providerID": ZCODE_PROVIDER,
                    "modelID": message.get("model") or plan["model"] or "unknown",
                },
            }
            for field in ("characterCount", "turnIndex"):
                if field in message and message[field] is not None:
                    message_data[field] = message[field]

        db.execute(insert_message, (
            message_id,
            plan["restoredTaskId"],
            message_timestamp,
            message_timestamp,
            jdumps(remove_undefined(message_data)),
        ))

        part_data_list = build_part_data_list(message, message_id, message_timestamp)
        for part_index, part_data in enumerate(part_data_list):
            part_timestamp = message_timestamp + part_index
            db.execute(insert_part, (
                f"part_legacy_{plan['restoredTaskId']}_"
                f"{str(index).zfill(4)}_{str(part_index).zfill(4)}",
                message_id,
                plan["restoredTaskId"],
                part_timestamp,
                part_timestamp,
                jdumps(part_data),
            ))
            part_count += 1

        if role == "user":
            latest_user_message_id = message_id

    return {"messageCount": len(plan["messages"]), "partCount": part_count}


def do_import(snapshot_path, task_index_path, cli_db_path, dry_run=False,
              backup=False):
    plan = build_restore_plan(snapshot_path)
    print_plan(plan, dry_run=dry_run)
    if dry_run:
        print("Dry run mode: no writes performed.\n")
        return

    ensure_existing_db(task_index_path, "Task index")
    ensure_existing_db(cli_db_path, "CLI DB")
    if backup:
        ti_backup = backup_db(task_index_path)
        cd_backup = backup_db(cli_db_path)
        print(f"Backed up task index: {ti_backup}")
        print(f"Backed up CLI DB:     {cd_backup}")

    ti = sqlite3.connect(task_index_path)
    cd = sqlite3.connect(cli_db_path)
    ti_in_tx = cd_in_tx = False
    try:
        ti.execute("BEGIN IMMEDIATE")
        ti_in_tx = True
        cd.execute("BEGIN IMMEDIATE")
        cd_in_tx = True

        print("[1/2] Writing task index...")
        write_task_index(ti, plan)
        print(f"  tasks row ready: {plan['restoredTaskId']}")

        print("[2/2] Writing CLI session DB...")
        write_session(cd, plan)
        result = write_messages(cd, plan)
        print(f"  message rows ready: {result['messageCount']}")
        print(f"  part rows ready:    {result['partCount']}")

        ti.commit()
        ti_in_tx = False
        cd.commit()
        cd_in_tx = False
        print("\n=== Import complete ===")
    except Exception as exc:
        if ti_in_tx:
            try:
                ti.execute("ROLLBACK")
            except Exception:
                pass
        if cd_in_tx:
            try:
                cd.execute("ROLLBACK")
            except Exception:
                pass
        print(f"Import failed: {exc}", file=sys.stderr)
        sys.exit(1)
    finally:
        ti.close()
        cd.close()


# ------------------------------------------------------------ official backend

def parse_version_string(value):
    out = []
    for part in re.split(r"[.\-+]", value):
        if part.isdigit():
            out.append(int(part))
        else:
            break
    return tuple(out)


def looks_like_scripts_dir(d):
    return (os.path.isfile(os.path.join(d, "scan-legacy-sessions.mjs"))
            and os.path.isfile(os.path.join(d, "restore-conversation.mjs")))


def resolve_scripts_dir():
    env = os.environ.get("ZCODE_LEGACY_SESSIONS_SCRIPTS")
    if env:
        d = os.path.expanduser(env)
        if looks_like_scripts_dir(d):
            return d
        raise SystemExit(
            f"ZCODE_LEGACY_SESSIONS_SCRIPTS set but no plugin scripts there: {d}")
    cache_root = os.path.join(HOME, ".zcode", "cli", "plugins", "cache",
                              "zcode-plugins-official",
                              "restore-legacy-sessions")
    if os.path.isdir(cache_root):
        for version in sorted(os.listdir(cache_root),
                              key=parse_version_string, reverse=True):
            d = os.path.join(cache_root, version, "skills",
                             "restore-legacy-sessions", "scripts")
            if looks_like_scripts_dir(d):
                return d
    bundled = os.path.join(
        os.environ.get("ProgramFiles", ""),
        "ZCode", "resources", "glm", "packages",
        "restore-legacy-sessions-plugin", "skills",
        "restore-legacy-sessions", "scripts")
    fallback = r"D:\ProgramFiles\ZCode\resources\glm\packages" \
               r"\restore-legacy-sessions-plugin\skills" \
               r"\restore-legacy-sessions\scripts"
    for base in (bundled, fallback):
        if base and looks_like_scripts_dir(base):
            return base
    return None


def detect_node():
    node = shutil.which("node")
    if not node:
        return None, []
    try:
        out = subprocess.run([node, "--version"], capture_output=True,
                             text=True, timeout=10)
        m = re.match(r"v(\d+)\.(\d+)", out.stdout.strip() or "")
    except Exception:
        return None, []
    if not m:
        return None, []
    major, minor = int(m.group(1)), int(m.group(2))
    if (major, minor) < (22, 5):
        return None, []
    flags = []
    if (major, minor) < (23, 4):
        flags = ["--experimental-sqlite"]
    return node, flags


def scan_filters(args, list_mode):
    extra = ["--legacy-dir", args["legacyDir"],
             "--task-index", args["taskIndex"],
             "--cli-db", args["cliDb"]]
    if args.get("agent") and args["agent"] != "all":
        extra += ["--agent", args["agent"]]
    if args.get("workspace"):
        extra += ["--workspace", args["workspace"]]
    if args.get("query"):
        extra += ["--query", args["query"]]
    if args.get("conversation"):
        extra += ["--conversation", args["conversation"]]
    if list_mode:
        extra += ["--limit", str(args.get("limit") or 30)]
    return extra


def run_scan(node, flags, scripts_dir, subcommand, extra):
    cmd = [node] + flags + [
        os.path.join(scripts_dir, "scan-legacy-sessions.mjs"),
        subcommand] + extra + ["--json"]
    proc = subprocess.run(cmd, capture_output=True, text=True,
                          encoding="utf-8", errors="replace")
    if proc.returncode != 0:
        detail = (proc.stderr.strip() or proc.stdout.strip() or "unknown error")
        raise SystemExit(f"scan {subcommand} failed (exit {proc.returncode}): "
                         f"{detail.splitlines()[-1]}")
    out = proc.stdout.strip()
    if not out:
        return None
    try:
        return json.loads(out)
    except Exception:
        raise SystemExit(f"scan {subcommand} returned non-JSON output")


def stream_official_import(node, flags, scripts_dir, snapshot, task_index,
                           cli_db, dry_run):
    cmd = [node] + flags + [
        os.path.join(scripts_dir, "restore-conversation.mjs"),
        "--snapshot", snapshot,
        "--task-index", task_index,
        "--cli-db", cli_db]
    if dry_run:
        cmd.append("--dry-run")
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT, bufsize=0)
    out_buf = getattr(sys.stdout, "buffer", None)
    try:
        while True:
            chunk = proc.stdout.read(65536)
            if not chunk:
                break
            if out_buf is not None:
                out_buf.write(chunk)
                out_buf.flush()
            else:
                sys.stdout.write(chunk.decode("utf-8", "replace"))
                sys.stdout.flush()
    finally:
        proc.wait()
        if proc.stdout:
            proc.stdout.close()
    return proc.returncode


def import_all_official(node, flags, scripts_dir, snapshots, task_index,
                        cli_db, dry_run):
    failed = 0
    total = len(snapshots)
    for i, snap in enumerate(snapshots, 1):
        print(f"[{i}/{total}] {snap}")
        sys.stdout.flush()
        code = stream_official_import(
            node, flags, scripts_dir, snap, task_index, cli_db, dry_run)
        if code != 0:
            failed += 1
            print(f"  (exit {code})", file=sys.stderr)
    if failed:
        raise SystemExit(f"import: {failed}/{total} failed")


def make_visible(task_index_path, task_ids):
    if not task_ids:
        return
    ensure_existing_db(task_index_path, "Task index")
    now = int(time.time() * 1000)
    placeholders = ",".join("?" for _ in task_ids)
    con = sqlite3.connect(task_index_path)
    try:
        con.execute("BEGIN IMMEDIATE")
        cur = con.execute(
            f"""UPDATE tasks
                SET archived = 0,
                    deleted = 0,
                    pinned = 0,
                    unread_at = ?
                WHERE task_id IN ({placeholders})""",
            [now] + list(task_ids))
        con.commit()
    except Exception:
        con.rollback()
        raise
    finally:
        con.close()
    print(f"made visible: {cur.rowcount} task(s) "
          f"(archived=0 deleted=0 pinned=0 unread_at={now})")


# ---------------------------------------------------------------------- CLI

def print_list(candidates, args):
    if args.get("json"):
        out = []
        for c in candidates:
            out.append({k: v for k, v in c.items() if k != "visibleText"})
        print(json.dumps(out, ensure_ascii=False, indent=2))
        return
    ordinals = build_global_ordinals(args["legacyDir"])
    rows = [[
        ordinals.get(os.path.abspath(c.get("filePath") or "")),
        c.get("restoreState") or "?",
        c["provider"],
        c["title"],
        c["messageCount"],
        ts(c.get("updatedAt") or 0),
        c["restoredTaskId"],
        c["legacyTaskId"],
    ] for c in sorted(
        candidates, key=lambda c: -(c.get("updatedAt") or 0))]
    print(md_table(
        ["#", "state", "agent", "title", "messages", "updated", "restoredTaskId", "legacyTaskId"],
        rows))


def print_status_view(by_agent, by_state, total, invalid, json_out=False):
    if json_out:
        print(json.dumps({
            "total": total,
            "invalid": invalid,
            "byAgent": by_agent,
            "byState": by_state,
        }, ensure_ascii=False, indent=2))
        return
    print(f"Total legacy conversations: {total}")
    if invalid:
        print(f"Invalid snapshots: {invalid}")
    print("\nBy agent:")
    print(md_table(
        ["agent", "conversations", "workspaces", "ready", "needs-cli-db",
         "needs-task-index", "needs-full-import"],
        [[r["key"], r["count"], r["workspaces"], r["ready"], r["needsCliDb"],
          r["needsTaskIndex"], r["needsFullImport"]] for r in by_agent]))
    print("\nBy restore state:")
    print(md_table(
        ["state", "conversations", "workspaces"],
        [[r["key"], r["count"], r["workspaces"]] for r in by_state]))


def print_status(candidates, json_out=False):
    valid = [c for c in candidates if not c.get("error")]
    by_agent = summarize_groups(valid, lambda c: c["provider"])
    by_state = summarize_groups(valid, lambda c: c["restoreState"])
    print_status_view(by_agent, by_state, len(valid),
                      len(candidates) - len(valid), json_out)


def parse_common(rest, args):
    leftovers = []
    i = 0
    while i < len(rest):
        a = rest[i]
        if a == "--legacy-dir":
            args["legacyDir"] = rest[i + 1]
            i += 2
        elif a in ("-t", "--task-index"):
            args["taskIndex"] = rest[i + 1]
            i += 2
        elif a in ("-c", "--cli-db"):
            args["cliDb"] = rest[i + 1]
            i += 2
        elif a == "--agent":
            args["agent"] = rest[i + 1]
            i += 2
        elif a == "--workspace":
            args["workspace"] = rest[i + 1]
            i += 2
        elif a == "--query":
            args["query"] = rest[i + 1]
            i += 2
        elif a == "--conversation":
            args["conversation"] = rest[i + 1]
            i += 2
        elif a == "--limit":
            args["limit"] = int(rest[i + 1])
            i += 2
        elif a == "--json":
            args["json"] = True
            i += 1
        elif a == "--dry-run":
            args["dryRun"] = True
            i += 1
        else:
            leftovers.append(a)
            i += 1
    return leftovers


def parse_import(rest, args):
    targets = []
    i = 0
    while i < len(rest):
        a = rest[i]
        if a == "--legacy-dir":
            args["legacyDir"] = rest[i + 1]
            i += 2
        elif a in ("-t", "--task-index"):
            args["taskIndex"] = rest[i + 1]
            i += 2
        elif a in ("-c", "--cli-db"):
            args["cliDb"] = rest[i + 1]
            i += 2
        elif a == "--dry-run":
            args["dryRun"] = True
            i += 1
        elif a == "--show":
            args["show"] = True
            i += 1
        elif a == "--no-show":
            args["show"] = False
            i += 1
        elif a == "--backup":
            args["backup"] = True
            i += 1
        elif a.startswith("--"):
            raise SystemExit(f"ERR: unknown import option: {a}")
        else:
            targets.append(a)
            i += 1
    if not targets:
        raise SystemExit("nothing to import (give snapshot paths, ordinals, a directory, or 'all')")
    if "all" in targets:
        targets = collect_json_files(args["legacyDir"])
    ordinals = build_global_ordinals(args["legacyDir"])
    file_by_ord = {v: k for k, v in ordinals.items()}
    expanded = []
    for t in targets:
        if isinstance(t, int) or (isinstance(t, str) and t.isdigit()):
            n = int(t)
            if n < 1:
                raise SystemExit(f"ordinal {n} out of range (1..{len(ordinals)})")
            fp = file_by_ord.get(n)
            if fp is None:
                raise SystemExit(
                    f"ordinal {n} out of range (1..{len(ordinals)})"
                    if ordinals else "nothing to import")
            expanded.append(fp)
        elif os.path.isdir(t):
            expanded.extend(collect_dir_snapshots(t))
        else:
            expanded.append(t)
    targets = expanded
    for t in targets:
        if not os.path.isfile(t):
            raise SystemExit(f"snapshot not found: {t}")
    return targets


def snapshot_restored_ids(snapshots):
    ids, unreadable = [], []
    for p in snapshots:
        c = read_snapshot(p)
        if c.get("error"):
            unreadable.append(p)
        elif c.get("restoredTaskId"):
            ids.append(c["restoredTaskId"])
    if unreadable:
        print(f"note: could not read restoredTaskId from "
              f"{len(unreadable)} snapshot(s)", file=sys.stderr)
    return ids


def resolve_task_ids(task_index_path, tasks):
    if tasks and "all" not in tasks:
        return list(tasks)
    ensure_existing_db(task_index_path, "Task index")
    con = sqlite3.connect(f"file:{task_index_path}?mode=ro", uri=True)
    try:
        return [r[0] for r in con.execute("SELECT task_id FROM tasks")]
    finally:
        con.close()


def main():
    argv = sys.argv[1:]
    if not argv or argv[0] not in ("list", "status", "import", "unarchive"):
        print(__doc__)
        sys.exit(1)
    cmd = argv[0]
    rest = argv[1:]
    args = {
        "legacyDir": LEGACY_DIR,
        "taskIndex": TASK_INDEX,
        "cliDb": CLI_DB,
        "json": False,
        "dryRun": False,
        "show": True,
        "backup": False,
        "limit": 30,
    }

    backend_env = (os.environ.get("ZCODE_IMPORT_BACKEND") or "official").strip().lower()
    if backend_env not in ("official", "python"):
        raise SystemExit(
            f"ZCODE_IMPORT_BACKEND must be 'official' or 'python': {backend_env}")
    run_official = backend_env == "official"

    if cmd == "import":
        targets = parse_import(rest, args)
        if run_official:
            scripts_dir = resolve_scripts_dir()
            node, flags = detect_node()
            if scripts_dir is not None and node is not None:
                before = set(list_dir_backup_files(args["taskIndex"]))
                before |= set(list_dir_backup_files(args["cliDb"]))
                import_all_official(node, flags, scripts_dir, targets,
                                    args["taskIndex"], args["cliDb"],
                                    args["dryRun"])
                if not args["backup"]:
                    cleanup_official_backups(
                        before, args["taskIndex"], args["cliDb"])
                if args["show"] and not args["dryRun"]:
                    make_visible(args["taskIndex"],
                                 snapshot_restored_ids(targets))
                return
            reason = "no node" if node is None else "official scripts not found"
            print(f"official backend unavailable ({reason}); "
                  f"using built-in Python port", file=sys.stderr)
        for t in targets:
            do_import(t, args["taskIndex"], args["cliDb"],
                      dry_run=args["dryRun"], backup=args.get("backup"))
        if args["show"] and not args["dryRun"]:
            make_visible(args["taskIndex"], snapshot_restored_ids(targets))
        return

    if cmd == "unarchive":
        leftovers = parse_common(rest, args)
        ids = resolve_task_ids(args["taskIndex"], leftovers or ["all"])
        make_visible(args["taskIndex"], ids)
        return

    leftovers = parse_common(rest, args)
    if leftovers:
        raise SystemExit(f"ERR: {cmd} takes options only")

    if run_official:
        scripts_dir = resolve_scripts_dir()
        node, flags = detect_node()
        if scripts_dir is not None and node is not None:
            if cmd == "list":
                rows = run_scan(node, flags, scripts_dir, "conversations",
                                scan_filters(args, list_mode=True))
                print_list(rows or [], args)
            else:
                payload = run_scan(node, flags, scripts_dir, "summary",
                                   scan_filters(args, list_mode=False)) or {}
                print_status_view(
                    payload.get("byAgent") or [],
                    payload.get("byState") or [],
                    int(payload.get("total") or 0),
                    int(payload.get("invalid") or 0),
                    args["json"])
            return
        reason = "no node" if node is None else "official scripts not found"
        print(f"official backend unavailable ({reason}); "
              f"using built-in Python port", file=sys.stderr)

    candidates = [read_snapshot(p) for p in collect_json_files(args["legacyDir"])]
    attach_store_status(candidates, args["taskIndex"], args["cliDb"])
    filtered = filter_candidates(candidates, args)
    if cmd == "list":
        if args.get("limit"):
            filtered = sorted(filtered, key=lambda c: -(c.get("updatedAt") or 0))
            filtered = filtered[:args["limit"]]
        print_list(filtered, args)
    else:
        print_status(filtered, args["json"])


if __name__ == "__main__":
    main()