opsxdb/src/store.ts

import { create } from "zustand";
import type { DriverId } from "./catalog";
import { driverById } from "./catalog";
import { formatDbError, notifyId, type NotifyKind, type NotifyPayload } from "./notify";
import { decodeSecret, encodeSecret } from "./secrets";
import { dropSchemaCache, syncConnectionsToDisk } from "./api";
import { normalizeRedisDatabase } from "./redisContext";

export type SavedConnection = {
  id: string;
  name: string;
  driver: DriverId;
  host: string;
  port: number;
  user: string;
  password: string;
  database: string;
  lastUsedAt?: number;
  /** 侧栏分组 id，缺省 default */
  groupId?: string;
  /** 组内排序，越小越靠前 */
  sortOrder?: number;
  ssl?: boolean;
  /** disable|require|skip-verify|verify-ca；关 SSL 时视为 disable */
  sslMode?: "disable" | "require" | "skip-verify" | "verify-ca";
  sslCA?: string;
  sslCert?: string;
  sslKey?: string;
  service?: string;
  color?: string;
  /** 连库页默认连接 */
  isDefault?: boolean;
  /** 是否持久化密码，默认 true */
  savePassword?: boolean;
  /** Redis key 分隔符，默认 : */
  keySeparator?: string;
  redisMode?: "standalone" | "sentinel" | "cluster";
  /** Sentinel/Cluster 节点列表 host:port（多行） */
  redisNodes?: string;
  /** Sentinel 主节点名 */
  masterName?: string;
  /** Sentinel 密码（可空） */
  sentinelPassword?: string;
  /** Oracle */
  oracleConnType?: "basic";
  oracleIdKind?: "service" | "sid";
  oracleRole?: string;
  /** SQL Server 验证方式 */
  authMode?: "sql" | "windows";
  /** SSH 隧道 */
  sshEnabled?: boolean;
  sshHost?: string;
  sshPort?: number;
  sshUser?: string;
  sshAuth?: "password" | "privateKey";
  sshPassword?: string;
  sshPrivateKey?: string;
  sshPassphrase?: string;
};

export const DEFAULT_CONN_GROUP = "default";

export type ConnectionGroup = {
  id: string;
  name: string;
};

export type QueryTab = {
  id: string;
  title: string;
  sql: string;
  connectionId?: string;
  database?: string;
  /** 相对打开/保存基线有改动 */
  dirty?: boolean;
  savedSql?: string;
};

export type QueryResult = {
  columns: string[];
  rows: Record<string, unknown>[];
  rowCount: number;
  timeMs: number;
  error?: string;
  truncated?: boolean;
  limit?: number;
  /** C5:服务端回传的分页偏移 */
  offset?: number;
  /** Redis 值游标分页：>0 还有更多（hash/set=SCAN 游标；list/zset=偏移） */
  nextCursor?: number;
  /** C1:命中危险语句,需确认后带 confirmDanger 重发 */
  needConfirm?: boolean;
};

export type SchemaView = {
  database: string;
  table: string;
  columns: ColumnInfo[];
  loading?: boolean;
  error?: string;
};

export type ResultMeta = {
  database?: string;
  table?: string;
  pkCols?: string[];
  /** Redis 键详情：打开键时带类型/TTL，命令结果不带 */
  redisType?: string;
  redisTtl?: number;
};

export type ConnState = "ok" | "down" | "reconnecting";

export type ResultSetItem = {
  id: string;
  label: string;
  sql: string;
  result: QueryResult;
  /** 仅对象预览填入，避免复制 INSERT 时误用其他标签的全局元数据。 */
  source?: { connectionId: string; driver: string; database: string; table: string };
  /** C5:分页状态(仅单条 SELECT 启用);page 从 0 起 */
  page?: number;
  pageSize?: number;
  /** 执行时使用的库(翻页重发需要) */
  database?: string;
  /** 列宽/选中态等视图持久化(修 G3,resultSet 级) */
  viewState?: {
    colWidths?: Record<string, number>;
    sortCol?: string | null;
    sortDir?: "asc" | "desc" | null;
    selectedIdx?: number | null;
    selectedCol?: string | null;
  };
};

export type Bookmark = {
  id: string;
  name: string;
  sql: string;
  at: number;
};

export type ColumnInfo = {
  name: string;
  type: string;
  nullable: boolean;
  key?: string;
  default?: unknown;
  comment?: string;
  extra?: string;
};

export type TreeNode = {
  name: string;
  kind: "database" | "group" | "table" | "view" | "procedure" | "function" | "column" | "key" | "folder";
  /** group id: tables|views|procedures|functions|keys|folder */
  groupId?: string;
  count?: number;
  children?: TreeNode[];
  loaded?: boolean;
  loading?: boolean;
  type?: string;
  /** Redis TTL seconds; -1 permanent; -2 missing */
  ttl?: number;
  /** Redis 文件夹前缀（含分隔符），如 CEX:FOLLOW: */
  fullName?: string;
  /** Redis TYPE 精确命中（搜索置顶标「精确」） */
  exact?: boolean;
  nullable?: boolean;
  key?: string;
};

export type QueryHistoryItem = {
  id: string;
  sql: string;
  at: number;
  connectionId?: string;
  rowCount?: number;
  timeMs?: number;
  error?: string;
};

export type ConfirmRequest = {
  title: string;
  message: string;
  okLabel?: string;
  danger?: boolean;
  /** C6:将执行的语句清单(危险确认二段流展示,前 3 条 + 折叠其余) */
  items?: string[];
};

export type AiAgentQueryResultView = {
  columns: string[];
  rows: Record<string, unknown>[];
  rowCount: number;
  sql?: string;
  tool?: string;
};

export type AiAgentStepView = {
  tool: string;
  args?: string;
  ok: boolean;
  summary: string;
  error?: string;
  ms: number;
  pending?: boolean;
  needConfirm?: boolean;
  confirmSql?: string;
  confirmBind?: string;
  result?: AiAgentQueryResultView;
};

export type AiMessage = {
  id: string;
  role: "user" | "assistant" | "system";
  content: string;
  at: string;
  runId?: string;
  sessionId?: string;
  runOutcome?: string;
  incompleteReason?: string;
  resumable?: boolean;
  stepsBudget?: number;
  /** Agent 工具步骤时间线（仅 assistant） */
  steps?: AiAgentStepView[];
  streaming?: boolean;
};

export type AiChatMode = "ask" | "agent";

export type ViewMode = "workbench" | "picker";
export type PickerMode = "list" | "drivers";

type State = {
  theme: "dark" | "light" | "system";
  view: ViewMode;
  connections: SavedConnection[];
  connGroups: ConnectionGroup[];
  expandedConnGroups: string[];
  activeConnectionId: string | null;
  tree: TreeNode[];
  treeLoading: boolean;
  /** 侧栏 ↻ 等触发对象树重载 */
  treeEpoch: number;
  tabs: QueryTab[];
  activeTabId: string;
  result: QueryResult | null;
  resultSets: ResultSetItem[];
  activeResultSetId: string | null;
  resultMeta: ResultMeta | null;
  schemaView: SchemaView | null;
  resultTab: "data" | "schema" | "explain";
  explainText: string | null;
  bookmarks: Bookmark[];
  showBookmarks: boolean;
  connState: ConnState;
  /** 左栏底部统一搜索：非 Redis=过滤库/表；Redis=过滤键（回车服务端模糊扫描） */
  sidebarFilter: string;
  /** Redis 类型筛选 */
  sidebarTypeFilter: string;
  running: boolean;
  pickerSearch: string;
  pickerCategory: string;
  pickerDriver: DriverId | null;
  pickerMode: PickerMode;
  form: Partial<SavedConnection>;
  aiMessages: AiMessage[];
  aiAsk: string;
  aiModel: string;
  aiBusy: boolean;
  /** Ask=不调工具；Agent=多轮 aitools */
  aiMode: AiChatMode;
  /** 当前 Agent 会话 id（落盘续跑） */
  aiSessionId: string | null;
  /** 忙时单槽排队文案；再提交覆盖 */
  aiQueuedPrompt: string | null;
  status: string;
  notifyModal: NotifyPayload | null;
  notifyToasts: NotifyPayload[];
  sessions: Record<string, string>;
  queryHistory: QueryHistoryItem[];
  leftWidth: number;
  aiWidth: number;
  resultRatio: number;
  showAi: boolean;
  showHelp: boolean;
  showHistory: boolean;
  confirm: ConfirmRequest | null;
  /** R2-F1:多语句执行遇错即停(默认 true);false 时全部执行完再汇总 */
  stopOnError: boolean;
  /** R2-F8:语句执行超时(ms),0 = 不超时;持久化到 prefs */
  execTimeoutMs: number;
  /** 对象树最近聚焦的表/视图(编辑器右键「查看数据/DDL」用) */
  focusedRelation: { database: string; table: string } | null;
  /** 请求 ObjectTree 打开 DDL 弹层 */
  ddlRequest: { database: string; table: string; token: number } | null;
  /** 请求 ObjectTree 打开表数据预览 */
  viewDataRequest: { database: string; table: string; token: number } | null;

  set: <K extends keyof State>(key: K, value: State[K]) => void;
  requestTreeRefresh: () => void;
  notify: (input: {
    kind?: NotifyKind;
    title?: string;
    message: string;
    detail?: string;
    tips?: string[];
    mode?: "modal" | "toast";
    /** 成功操作默认静默；长任务完成等需要用户注意时显式设为 visible。 */
    feedback?: "visible";
    durationMs?: number;
    raw?: unknown;
  }) => void;
  notifyError: (raw: unknown, mode?: "modal" | "toast") => void;
  dismissNotify: (id: string) => void;
  loadConnections: () => void;
  saveConnections: (list: SavedConnection[]) => void;
  upsertConnection: (c: SavedConnection) => void;
  removeConnection: (id: string) => void;
  setDefaultConnection: (id: string) => void;
  addConnGroup: (name: string) => string;
  /** R2-F5:导入连接时并入分组(冲突/缺失 id 重新生成);默认分组不重复创建 */
  importConnGroups: (groups: ConnectionGroup[]) => void;
  renameConnGroup: (id: string, name: string) => void;
  removeConnGroup: (id: string) => void;
  moveConnection: (id: string, groupId: string, index: number) => void;
  toggleConnGroup: (id: string) => void;
  openPicker: (driver?: DriverId | null) => void;
  openConnectionList: () => void;
  editConnection: (c: SavedConnection) => void;
  newQueryTab: (database?: string) => void;
  /** 重命名当前激活的查询 Tab（如打开表后标题设为表名） */
  setActiveTabTitle: (title: string) => void;
  /** 设置当前 Tab 的库上下文（Redis 编辑器 SELECT/USE 切库用） */
  setTabDatabase: (database: string, tabId?: string) => void;
  /** 切连接时激活其查询标签，或新建查询；保留别的连接的 SQL 草稿。 */
  selectConnection: (id: string) => void;
  closeTab: (id: string) => void;
  /** R2-F6:Tab 右键菜单批量关闭(脏确认由调用侧负责) */
  closeOtherTabs: (id: string) => void;
  closeRightTabs: (id: string) => void;
  closeAllTabs: () => void;
  setStopOnError: (v: boolean) => void;
  setExecTimeoutMs: (v: number) => void;
  updateActiveSql: (sql: string) => void;
  /** AI 异步返回只写入原标签；用户编辑或切库后跳过。 */
  updateTabSqlIfUnchanged: (id: string, originalSql: string, connectionId: string | undefined, database: string | undefined, sql: string) => boolean;
  pushAi: (role: AiMessage["role"], content: string, steps?: AiAgentStepView[], extra?: Partial<AiMessage>) => string;
  patchAiMessage: (id: string, patch: Partial<AiMessage> | ((m: AiMessage) => AiMessage)) => void;
  setSession: (connectionId: string, sessionId: string | null) => void;
  pushHistory: (item: Omit<QueryHistoryItem, "id" | "at"> & { at?: number }) => void;
  clearHistory: () => void;
  setResultSets: (items: ResultSetItem[]) => void;
  /** AI「在结果区打开」：新开/切换结果页，不改 SQL Tab、不抢编辑器 focus */
  openAiResultSet: (label: string, sql: string, result: QueryResult) => void;
  selectResultSet: (id: string) => void;
  /** C5/G3:局部更新某个结果集(翻页结果、视图持久化);若为激活项同步 result */
  updateResultSet: (id: string, patch: Partial<Omit<ResultSetItem, "id">>) => void;
  loadBookmarks: () => void;
  saveBookmark: (name: string, sql: string) => void;
  renameBookmark: (id: string, name: string) => void;
  removeBookmark: (id: string) => void;
  askConfirm: (req: ConfirmRequest) => Promise<boolean>;
  resolveConfirm: (ok: boolean) => void;
  persistPrefs: () => void;
};

const LS_KEY = "opsxdb.connections.v1";
const LS_GROUPS = "opsxdb.conn-groups.v1";
const LS_EXPANDED = "opsxdb.conn-groups-expanded.v1";
const LS_PREFS = "opsxdb.prefs.v1";
const LS_TABS = "opsxdb.tabs.v1";
const LS_HISTORY = "opsxdb.history.v1";
const LS_BOOKMARKS = "opsxdb.bookmarks.v1";

let confirmResolver: ((ok: boolean) => void) | null = null;

function loadPrefs(): Partial<Pick<State, "theme" | "leftWidth" | "aiWidth" | "resultRatio" | "showAi" | "stopOnError" | "execTimeoutMs">> {
  try {
    const raw = localStorage.getItem(LS_PREFS);
    if (!raw) return {};
    const p = JSON.parse(raw) as Record<string, unknown>;
    return {
      theme: p.theme === "light" || p.theme === "system" ? p.theme : "dark",
      leftWidth: typeof p.leftWidth === "number" ? Math.min(360, Math.max(200, p.leftWidth)) : undefined,
      aiWidth: typeof p.aiWidth === "number" ? p.aiWidth : undefined,
      resultRatio: typeof p.resultRatio === "number" ? p.resultRatio : undefined,
      showAi: typeof p.showAi === "boolean" ? p.showAi : undefined,
      stopOnError: typeof p.stopOnError === "boolean" ? p.stopOnError : undefined,
      execTimeoutMs: typeof p.execTimeoutMs === "number" ? Math.min(3_600_000, Math.max(0, p.execTimeoutMs)) : undefined,
    };
  } catch {
    return {};
  }
}

function loadTabs(): { tabs: QueryTab[]; activeTabId: string } | null {
  try {
    const raw = localStorage.getItem(LS_TABS);
    if (!raw) return null;
    const p = JSON.parse(raw) as { tabs?: QueryTab[]; activeTabId?: string };
    if (!Array.isArray(p.tabs) || !p.tabs.length) return null;
    const tabs = p.tabs.map((tab) => {
      const savedSql = tab.savedSql ?? tab.sql ?? "";
      const sql = tab.sql ?? "";
      return { ...tab, sql, savedSql, dirty: sql !== savedSql };
    });
    return { tabs, activeTabId: p.activeTabId || tabs[0].id };
  } catch {
    return null;
  }
}

function loadHistory(): QueryHistoryItem[] {
  try {
    const raw = localStorage.getItem(LS_HISTORY);
    if (!raw) return [];
    const p = JSON.parse(raw);
    return Array.isArray(p) ? p.slice(0, 50) : [];
  } catch {
    return [];
  }
}

function loadGroups(): ConnectionGroup[] {
  try {
    const raw = localStorage.getItem(LS_GROUPS);
    if (!raw) return [{ id: DEFAULT_CONN_GROUP, name: "默认" }];
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed) || !parsed.length) return [{ id: DEFAULT_CONN_GROUP, name: "默认" }];
    const list = parsed.filter((g: ConnectionGroup) => g && g.id && g.name);
    if (!list.some((g: ConnectionGroup) => g.id === DEFAULT_CONN_GROUP)) {
      list.unshift({ id: DEFAULT_CONN_GROUP, name: "默认" });
    }
    return list;
  } catch {
    return [{ id: DEFAULT_CONN_GROUP, name: "默认" }];
  }
}

function saveGroups(list: ConnectionGroup[]) {
  localStorage.setItem(LS_GROUPS, JSON.stringify(list));
}

function loadExpanded(): string[] {
  try {
    const raw = localStorage.getItem(LS_EXPANDED);
    if (!raw) return [DEFAULT_CONN_GROUP];
    const parsed = JSON.parse(raw);
    return Array.isArray(parsed) ? parsed : [DEFAULT_CONN_GROUP];
  } catch {
    return [DEFAULT_CONN_GROUP];
  }
}

function loadLs(): SavedConnection[] {
  try {
    const raw = localStorage.getItem(LS_KEY);
    if (!raw) return [];
    const parsed = JSON.parse(raw);
    return Array.isArray(parsed) ? parsed : [];
  } catch {
    return [];
  }
}

function defaultTab(): QueryTab {
  // 新 Tab 默认空白(修 T3),不再预置 SELECT 1 AS ok;
  return { id: crypto.randomUUID(), title: "查询 1", sql: "", savedSql: "", dirty: false };
}

/** 启动占位标签没有任何用户内容，可在真正创建标签时直接替换。 */
function isEmptyStartupTab(tab: QueryTab): boolean {
  return !tab.connectionId && !tab.database && !tab.sql?.trim() && !tab.savedSql?.trim() && /^查询 1$/.test(tab.title);
}

export const useWorkspaceStore = create<State>((set, get) => ({
  theme: "dark",
  view: "workbench",
  connections: [],
  connGroups: [{ id: DEFAULT_CONN_GROUP, name: "默认" }],
  expandedConnGroups: [DEFAULT_CONN_GROUP],
  activeConnectionId: null,
  tree: [],
  treeLoading: false,
  treeEpoch: 0,
  tabs: [defaultTab()],
  activeTabId: "",
  result: null,
  resultSets: [],
  activeResultSetId: null,
  resultMeta: null,
  schemaView: null,
  resultTab: "data",
  explainText: null,
  bookmarks: [],
  showBookmarks: false,
  connState: "ok",
  sidebarFilter: "",
  sidebarTypeFilter: "all",
  running: false,
  pickerSearch: "",
  pickerCategory: "全部",
  pickerDriver: null,
  pickerMode: "list",
  form: {},
  aiMessages: [],
  aiAsk: "",
  aiModel: "opsxdb-stub",
  aiBusy: false,
  aiMode: "ask",
  aiSessionId: null,
  aiQueuedPrompt: null,
  status: "就绪",
  notifyModal: null,
  notifyToasts: [],
  sessions: {},
  queryHistory: [],
  leftWidth: 260,
  aiWidth: 280,
  resultRatio: 0.38,
  showAi: true,
  showHelp: false,
  showHistory: false,
  confirm: null,
  stopOnError: true,
  execTimeoutMs: 0,
  focusedRelation: null,
  ddlRequest: null,
  viewDataRequest: null,

  set: (key, value) => {
    set((state) => key === "activeConnectionId" && value !== state.activeConnectionId
      ? ({ [key]: value, sidebarFilter: "" } as Partial<State>)
      : ({ [key]: value } as Partial<State>));
    if (key === "theme" || key === "leftWidth" || key === "aiWidth" || key === "resultRatio" || key === "showAi" || key === "stopOnError" || key === "execTimeoutMs") {
      get().persistPrefs();
    }
    if (key === "tabs" || key === "activeTabId") {
      const s = get();
      localStorage.setItem(LS_TABS, JSON.stringify({ tabs: s.tabs, activeTabId: s.activeTabId }));
    }
  },

  requestTreeRefresh: () => set((st) => ({ treeEpoch: (st.treeEpoch || 0) + 1 })),

  notify: (input) => {
    const parsed = input.raw !== undefined ? formatDbError(input.raw) : null;
    const kind = input.kind ?? parsed?.kind ?? "info";
    const mode = input.mode ?? (kind === "error" ? "modal" : "toast");
    if (kind === "success" && mode === "toast" && input.feedback !== "visible") return;
    const payload: NotifyPayload = {
      id: notifyId(),
      kind,
      title: input.title ?? parsed?.title ?? (kind === "error" ? "错误" : "提示"),
      message: input.message || parsed?.message || "",
      detail: input.detail ?? parsed?.detail,
      tips: input.tips ?? parsed?.tips,
      mode,
      durationMs: input.durationMs,
    };
    if (mode === "modal") {
      set({ notifyModal: payload });
    } else {
      set((s) => ({ notifyToasts: [...s.notifyToasts.slice(-4), payload] }));
    }
  },

  notifyError: (raw, mode = "modal") => {
    const parsed = formatDbError(raw);
    const payload: NotifyPayload = {
      id: notifyId(),
      kind: parsed.kind,
      title: parsed.title,
      message: parsed.message,
      detail: parsed.detail,
      tips: parsed.tips,
      mode,
    };
    if (mode === "modal") set({ notifyModal: payload });
    else set((s) => ({ notifyToasts: [...s.notifyToasts.slice(-4), payload] }));
  },

  dismissNotify: (id) => {
    const s = get();
    if (s.notifyModal?.id === id) set({ notifyModal: null });
    set({ notifyToasts: get().notifyToasts.filter((t) => t.id !== id) });
  },

  loadConnections: () => {
    const list = loadLs().map((c) => ({
      ...c,
      groupId: c.groupId || DEFAULT_CONN_GROUP,
      sortOrder: c.sortOrder ?? 0,
    }));
    const prefs = loadPrefs();
    const savedTabs = loadTabs();
    const loadedTabs = savedTabs?.tabs ?? get().tabs;
    const tabs = loadedTabs.length > 1 ? loadedTabs.filter(t => !isEmptyStartupTab(t)) : loadedTabs;
    const activeTabId = tabs.some(t => t.id === savedTabs?.activeTabId)
      ? savedTabs!.activeTabId : tabs[0]?.id || "";
    set({
      connections: list,
      connGroups: loadGroups(),
      expandedConnGroups: loadExpanded(),
      tabs,
      activeTabId,
      queryHistory: loadHistory(),
      ...prefs,
    });
  },

  saveConnections: (list) => {
    const encoded = list.map((c) => ({
      ...c,
      password: encodeSecret(c.password || ""),
      sshPassword: c.sshPassword ? encodeSecret(c.sshPassword) : c.sshPassword,
      sshPassphrase: c.sshPassphrase ? encodeSecret(c.sshPassphrase) : c.sshPassphrase,
      sentinelPassword: c.sentinelPassword ? encodeSecret(c.sentinelPassword) : c.sentinelPassword,
    }));
    localStorage.setItem(LS_KEY, JSON.stringify(encoded));
    set({ connections: encoded });
    // 双写 ~/.opsxdb/connections.json（经本机 Gin），供 opsxdb-mcp 读取。 // Logan.Li
    void syncConnectionsToDisk(encoded);
  },

  upsertConnection: (c) => {
    const next = {
      ...c,
      groupId: c.groupId || DEFAULT_CONN_GROUP,
      sortOrder: c.sortOrder ?? Date.now(),
      password: c.savePassword === false ? "" : c.password,
      sshPassword: c.savePassword === false ? "" : c.sshPassword,
    };
    let list = [...get().connections];
    // 同名防护：不同 id 同名时自动加序号后缀，避免侧栏出现重名条目
    const nameTaken = (n: string) => list.some((x) => x.id !== next.id && x.name === n);
    if (nameTaken(next.name)) {
      let k = 2;
      while (nameTaken(`${next.name} ${k}`)) k++;
      next.name = `${next.name} ${k}`;
    }
    if (next.isDefault) {
      list = list.map((x) => ({ ...x, isDefault: x.id === next.id ? true : false }));
    }
    const i = list.findIndex((x) => x.id === next.id);
    if (i >= 0) list[i] = { ...list[i], ...next };
    else list.unshift(next);
    get().saveConnections(list);
  },

  removeConnection: (id) => {
    get().saveConnections(get().connections.filter((c) => c.id !== id));
    dropSchemaCache(id);
    const sessions = { ...get().sessions };
    delete sessions[id];
    if (get().activeConnectionId === id) set({ activeConnectionId: null, sidebarFilter: "", tree: [], sessions });
    else set({ sessions });
  },

  setDefaultConnection: (id) => {
    const list = get().connections.map((c) => ({ ...c, isDefault: c.id === id }));
    get().saveConnections(list);
  },


  addConnGroup: (name) => {
    const n = name.trim() || "新分组";
    const id = crypto.randomUUID();
    const list = [...get().connGroups, { id, name: n }];
    saveGroups(list);
    const expanded = Array.from(new Set([...get().expandedConnGroups, id]));
    localStorage.setItem(LS_EXPANDED, JSON.stringify(expanded));
    set({ connGroups: list, expandedConnGroups: expanded });
    return id;
  },

  /** R2-F5:导入连接时并入分组;与现有 id 冲突或缺失时重新生成,避免串组 */
  importConnGroups: (groups) => {
    if (!Array.isArray(groups) || !groups.length) return;
    const cur = get().connGroups;
    const ids = new Set(cur.map((g) => g.id));
    const merged = [...cur];
    for (const g of groups) {
      if (!g || typeof g.name !== "string" || !g.name.trim()) continue;
      if (g.id === DEFAULT_CONN_GROUP) continue;
      const item = { id: g.id && !ids.has(g.id) ? g.id : crypto.randomUUID(), name: g.name.trim() };
      ids.add(item.id);
      merged.push(item);
    }
    saveGroups(merged);
    set({ connGroups: merged });
  },

  renameConnGroup: (id, name) => {
    if (id === DEFAULT_CONN_GROUP) return;
    const n = name.trim();
    if (!n) return;
    const list = get().connGroups.map((g) => (g.id === id ? { ...g, name: n } : g));
    saveGroups(list);
    set({ connGroups: list });
  },

  removeConnGroup: (id) => {
    if (id === DEFAULT_CONN_GROUP) return;
    const list = get().connGroups.filter((g) => g.id !== id);
    saveGroups(list);
    const connections = get().connections.map((c) =>
      (c.groupId || DEFAULT_CONN_GROUP) === id ? { ...c, groupId: DEFAULT_CONN_GROUP } : c,
    );
    get().saveConnections(connections);
    const expanded = get().expandedConnGroups.filter((x) => x !== id);
    localStorage.setItem(LS_EXPANDED, JSON.stringify(expanded));
    set({ connGroups: list, expandedConnGroups: expanded });
  },

  moveConnection: (id, groupId, index) => {
    const gid = groupId || DEFAULT_CONN_GROUP;
    const all = [...get().connections];
    const moving = all.find((c) => c.id === id);
    if (!moving) return;
    const othersInGroup = all
      .filter((c) => c.id !== id && (c.groupId || DEFAULT_CONN_GROUP) === gid)
      .sort((a, b) => (a.sortOrder ?? 0) - (b.sortOrder ?? 0));
    const clamped = Math.max(0, Math.min(index, othersInGroup.length));
    othersInGroup.splice(clamped, 0, { ...moving, groupId: gid });
    othersInGroup.forEach((c, i) => {
      c.sortOrder = i;
    });
    const rest = all.filter((c) => c.id !== id && (c.groupId || DEFAULT_CONN_GROUP) !== gid);
    get().saveConnections([...rest, ...othersInGroup]);
    if (!get().expandedConnGroups.includes(gid)) {
      const expanded = [...get().expandedConnGroups, gid];
      localStorage.setItem(LS_EXPANDED, JSON.stringify(expanded));
      set({ expandedConnGroups: expanded });
    }
  },

  toggleConnGroup: (id) => {
    const cur = get().expandedConnGroups;
    const expanded = cur.includes(id) ? cur.filter((x) => x !== id) : [...cur, id];
    localStorage.setItem(LS_EXPANDED, JSON.stringify(expanded));
    set({ expandedConnGroups: expanded });
  },

  openPicker: (driver = null) => {
    const d = driver ? driverById(driver) : null;
    set({
      view: "picker",
      pickerMode: d ? "drivers" : "drivers",
      pickerDriver: driver,
      form: d
        ? {
            id: crypto.randomUUID(),
            name: `${d.name} 连接`,
            driver: d.id,
            host: d.hints?.host ?? "127.0.0.1",
            port: d.defaultPort,
            user: d.hints?.user ?? (d.id === "redis" ? "" : "root"),
            password: "",
            database: d.hints?.database ?? "",
            ssl: false,
            sslMode: "disable",
            sslCA: "",
            sslCert: "",
            sslKey: "",
            service: d.id === "oracle" ? (d.hints?.database ?? "ORCL") : "",
            savePassword: true,
            keySeparator: d.id === "redis" ? ":" : undefined,
            redisMode: d.id === "redis" ? "standalone" : undefined,
            oracleConnType: d.id === "oracle" ? "basic" : undefined,
            oracleIdKind: d.id === "oracle" ? "service" : undefined,
            oracleRole: d.id === "oracle" ? "default" : undefined,
            authMode: d.id === "sqlserver" ? "sql" : undefined,
            sshEnabled: false,
            sshPort: 22,
            sshAuth: "password",
          }
        : {},
    });
  },

  openConnectionList: () => {
    set({ view: "picker", pickerMode: "list", pickerDriver: null, form: {} });
  },

  editConnection: (c) => {
    set({
      view: "picker",
      pickerMode: "drivers",
      pickerDriver: c.driver,
      form: {
        ...c,
        password: decodeSecret(c.password),
        sshPassword: c.sshPassword ? decodeSecret(c.sshPassword) : "",
        sshPassphrase: c.sshPassphrase ? decodeSecret(c.sshPassphrase) : "",
      },
    });
  },

  newQueryTab: (database) => {
    const current = get();
    const existingTabs = current.tabs.filter(t => !isEmptyStartupTab(t));
    const n = existingTabs.length + 1;
    const conn = get().connections.find((c) => c.id === get().activeConnectionId);
    const requestedDatabase = database ?? conn?.database;
    const tab: QueryTab = {
      id: crypto.randomUUID(),
      title: `查询 ${n}`,
      sql: "",
      savedSql: "",
      dirty: false,
      connectionId: get().activeConnectionId ?? undefined,
      database: conn?.driver.toLowerCase() === "redis" ? normalizeRedisDatabase(requestedDatabase, conn) : requestedDatabase,
    };
    const tabs = [...existingTabs, tab];
    set({ tabs, activeTabId: tab.id });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: tab.id }));
  },

  setActiveTabTitle: (title) => {
    const s = get();
    const tabs = s.tabs.map((t) => (t.id === s.activeTabId ? { ...t, title } : t));
    set({ tabs });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: s.activeTabId }));
  },

  setTabDatabase: (database, tabId) => {
    const s = get();
    const id = tabId ?? s.activeTabId;
    const target = s.tabs.find(t => t.id === id);
    if (!target) return;
    const conn = s.connections.find((c) => c.id === target.connectionId);
    const nextDatabase = conn?.driver.toLowerCase() === "redis" ? normalizeRedisDatabase(database, conn) : database;
    const tabs = s.tabs.map((t) => (t.id === id ? { ...t, database: nextDatabase } : t));
    set({ tabs });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: s.activeTabId }));
  },

  selectConnection: (id) => {
    if (!get().connections.some(c => c.id === id)) return;
    get().set("activeConnectionId", id);
    const s = get();
    const conn = s.connections.find(c => c.id === id)!;
    const normalizeTab = (tabId: string) => {
      if (conn.driver.toLowerCase() !== "redis") return;
      const current = get();
      const tabs = current.tabs.map(t => t.id === tabId ? { ...t, database: normalizeRedisDatabase(t.database ?? conn.database, conn) } : t);
      set({ tabs });
      localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: current.activeTabId }));
    };
    if (s.tabs.find(t => t.id === s.activeTabId)?.connectionId === id) {
      normalizeTab(s.activeTabId);
      return;
    }
    const existing = s.tabs.find(t => t.connectionId === id);
    if (existing) {
      get().set("activeTabId", existing.id);
      normalizeTab(existing.id);
    }
    else get().newQueryTab();
  },

  closeTab: (id) => {
    const tabs = get().tabs.filter((t) => t.id !== id);
    const next = tabs.length ? tabs : [defaultTab()];
    const activeTabId = get().activeTabId === id ? next[0].id : get().activeTabId;
    set({ tabs: next, activeTabId });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs: next, activeTabId }));
  },

  /** R2-F6:右键菜单「关闭其他」——只保留指定 Tab */
  closeOtherTabs: (id) => {
    const tabs = get().tabs.filter((t) => t.id === id);
    const next = tabs.length ? tabs : [defaultTab()];
    const activeTabId = next.some((t) => t.id === get().activeTabId) ? get().activeTabId : next[0].id;
    set({ tabs: next, activeTabId });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs: next, activeTabId }));
  },

  /** R2-F6:右键菜单「关闭右侧」——保留指定 Tab 及其左侧全部 */
  closeRightTabs: (id) => {
    const idx = get().tabs.findIndex((t) => t.id === id);
    if (idx < 0) return;
    const next = get().tabs.slice(0, idx + 1);
    const activeTabId = next.some((t) => t.id === get().activeTabId) ? get().activeTabId : next[next.length - 1].id;
    set({ tabs: next, activeTabId });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs: next, activeTabId }));
  },

  /** R2-F6:右键菜单「关闭全部」——回落一个空白 Tab(与 closeTab 语义一致) */
  closeAllTabs: () => {
    const next = [defaultTab()];
    set({ tabs: next, activeTabId: next[0].id });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs: next, activeTabId: next[0].id }));
  },

  setStopOnError: (v) => {
    set({ stopOnError: v });
    get().persistPrefs();
  },

  setExecTimeoutMs: (v) => {
    set({ execTimeoutMs: Math.min(3_600_000, Math.max(0, v || 0)) });
    get().persistPrefs();
  },

  updateActiveSql: (sql) => {
    const id = get().activeTabId;
    const tabs = get().tabs.map((t) => {
      if (t.id !== id) return t;
      const saved = t.savedSql ?? "";
      return { ...t, sql, dirty: sql !== saved };
    });
    set({ tabs });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: id }));
  },

  updateTabSqlIfUnchanged: (id, originalSql, connectionId, database, sql) => {
    const current = get().tabs.find(t => t.id === id);
    if (!current || current.sql !== originalSql || current.connectionId !== connectionId || current.database !== database) return false;
    const tabs = get().tabs.map(t => t.id === id ? { ...t, sql, dirty: sql !== (t.savedSql ?? "") } : t);
    set({ tabs });
    localStorage.setItem(LS_TABS, JSON.stringify({ tabs, activeTabId: get().activeTabId }));
    return true;
  },

  pushAi: (role, content, steps, extra) => {
    const id = crypto.randomUUID();
    set((s) => ({
      aiMessages: [
        ...s.aiMessages,
        {
          id,
          role,
          content,
          at: new Date().toLocaleTimeString(),
          steps,
          ...extra,
        },
      ],
    }));
    return id;
  },

  patchAiMessage: (id, patch) =>
    set((s) => ({
      aiMessages: s.aiMessages.map((m) => {
        if (m.id !== id) return m;
        return typeof patch === "function" ? patch(m) : { ...m, ...patch };
      }),
    })),


  setSession: (connectionId, sessionId) => {
    const sessions = { ...get().sessions };
    if (sessionId) sessions[connectionId] = sessionId;
    else delete sessions[connectionId];
    set({ sessions });
  },

  pushHistory: (item) => {
    const norm = (s: string) => s.replace(/\s+/g, " ").trim().toLowerCase();
    const key = norm(item.sql);
    const entry: QueryHistoryItem = {
      id: crypto.randomUUID(),
      at: item.at ?? Date.now(),
      sql: item.sql,
      connectionId: item.connectionId,
      rowCount: item.rowCount,
      timeMs: item.timeMs,
      error: item.error,
    };
    // 相同语句去重置顶：重复执行只留最新一次，历史即「最近用过的不同语句」
    const next = [entry, ...get().queryHistory.filter((h) => norm(h.sql) !== key)].slice(0, 200);
    localStorage.setItem(LS_HISTORY, JSON.stringify(next));
    set({ queryHistory: next });
  },

  clearHistory: () => {
    set({ queryHistory: [] });
    localStorage.setItem(LS_HISTORY, JSON.stringify([]));
  },


  setResultSets: (items) => {
    const active = items[0]?.id ?? null;
    set({
      resultSets: items,
      activeResultSetId: active,
      result: items[0]?.result ?? null,
      resultTab: "data",
    });
  },
  openAiResultSet: (label, sql, result) => {
    const id = crypto.randomUUID();
    const item: ResultSetItem = { id, label: label || "AI 结果", sql: sql || "", result };
    const prev = get().resultSets.filter((x) => x.id !== id);
    set({
      resultSets: [item, ...prev].slice(0, 30),
      activeResultSetId: id,
      result,
      resultTab: "data",
    });
  },


  selectResultSet: (id) => {
    const item = get().resultSets.find((x) => x.id === id);
    if (!item) return;
    set({ activeResultSetId: id, result: item.result, resultTab: "data" });
  },

  updateResultSet: (id, patch) => {
    const items = get().resultSets.map((x) => (x.id === id ? { ...x, ...patch } : x));
    const next: Partial<State> = { resultSets: items };
    if (get().activeResultSetId === id) {
      const item = items.find((x) => x.id === id);
      if (item) next.result = item.result;
    }
    set(next);
  },

  loadBookmarks: () => {
    try {
      const raw = localStorage.getItem(LS_BOOKMARKS);
      if (!raw) return;
      const parsed = JSON.parse(raw);
      if (Array.isArray(parsed)) set({ bookmarks: parsed });
    } catch {
      /* ignore */
    }
  },

  saveBookmark: (name, sql) => {
    const n = name.trim();
    if (!n) return;
    const list = [{ id: crypto.randomUUID(), name: n, sql, at: Date.now() }, ...get().bookmarks].slice(0, 100);
    localStorage.setItem(LS_BOOKMARKS, JSON.stringify(list));
    set({ bookmarks: list });
  },

  renameBookmark: (id, name) => {
    const n = name.trim();
    if (!n) return;
    const list = get().bookmarks.map((b) => (b.id === id ? { ...b, name: n } : b));
    localStorage.setItem(LS_BOOKMARKS, JSON.stringify(list));
    set({ bookmarks: list });
  },

  removeBookmark: (id) => {
    const list = get().bookmarks.filter((b) => b.id !== id);
    localStorage.setItem(LS_BOOKMARKS, JSON.stringify(list));
    set({ bookmarks: list });
  },

  askConfirm: (req) =>
    new Promise<boolean>((resolve) => {
      confirmResolver?.(false);
      confirmResolver = resolve;
      set({ confirm: req });
    }),

  resolveConfirm: (ok) => {
    const r = confirmResolver;
    confirmResolver = null;
    set({ confirm: null });
    r?.(ok);
  },

  persistPrefs: () => {
    const s = get();
    localStorage.setItem(
      LS_PREFS,
      JSON.stringify({
        theme: s.theme,
        leftWidth: s.leftWidth,
        aiWidth: s.aiWidth,
        resultRatio: s.resultRatio,
        showAi: s.showAi,
        stopOnError: s.stopOnError,
        execTimeoutMs: s.execTimeoutMs,
      }),
    );
  },
}));

// init bookmarks
try {
  const _bm = localStorage.getItem(LS_BOOKMARKS);
  if (_bm) {
    const parsed = JSON.parse(_bm);
    if (Array.isArray(parsed)) useWorkspaceStore.setState({ bookmarks: parsed });
  }
} catch { /* ignore */ }

// init activeTabId after create
const boot = useWorkspaceStore.getState();
if (!boot.activeTabId && boot.tabs[0]) {
  useWorkspaceStore.setState({ activeTabId: boot.tabs[0].id });
}