import { create } from "zustand";
import { persist } from "zustand/middleware";
import type {
  AdminChannelMessage,
  AgentInfo,
  ApprovalNotice,
  AwakenState,
  BackthinkItem,
  ChatMessage,
  FlowGraph,
  Mode,
  PersonaRecord,
  SessionMeta,
  ToolUsed,
  WsStatus,
} from "../lib/types";
import { now, uid } from "../lib/util";
import {
  currentGraph,
  dropMessage,
  mergeHistory,
  mergeMessage,
  patchAssistantTail,
  patchMessage,
  writeFlowGraph,
  replacePage,
  withBucket,
  writeSession,
  writeSessionById,
  type SessionMessages,
} from "./sessionCache";
import { deliveryActions, turnActions, turnState, type TurnSlice } from "./turnSlice";
import { patchAssistantForTurn } from "./turnCache";
import { persistOptions } from "./persistOptions";

interface ChatContext {
  topic?: string;
  lastAgentMode?: string;
  notes?: string[];
}

export interface AppStore extends TurnSlice {
  /** Session metadata (main + secondary) keyed by session id. The message truth
   *  lives server-side, so this is a small list — safe to persist. */
  sessions: Record<string, SessionMeta>;
  /** The session the UI is currently focused on (design/16 §2.1). "" until the
   *  session list is loaded. Messages in `messages` always belong to this id. */
  activeSessionId: string;
  wsStatus: WsStatus;
  /** Messages of the *active* session, in display order (oldest first). Switching
   *  sessions no longer clears it: the list on screen is stashed into that session's
   *  bucket in `sessionMessages` and the target session's bucket is taken out
   *  (design/16 §4.3; 裁决 16.4.5). */
  messages: ChatMessage[];
  /** Message buckets of the sessions that are **not** on screen — the active
   *  session's list lives in `messages`. A frame (persisted message, stream delta,
   *  edit, delete) is written into the bucket its `sessionId` names instead of being
   *  dropped, so a background conversation keeps accumulating and is rendered whole
   *  the moment it is switched to (裁决 16.4.5). Bounded by what the backend pushed
   *  during this page's lifetime; the backend stays the message truth (design/16 §1)
   *  and a switch-in is refetched page by page. Never persisted. */
  sessionMessages: SessionMessages;
  /** Latest mood per role id (design/22 §4.6). The `mood` WS event is
   *  persona-scoped and carries no `sessionId`, so each role gets its own bucket
   *  and the UI shows the bucket of the *active session's* role — role X's mood
   *  must not sit above role Y's conversation. `""` holds events that arrived
   *  without a persona (legacy/unscoped frames). Seeded from the role catalog's
   *  role-level mood on load (`seedMoods` + `personaMoodSeed`). */
  moods: Record<string, string>;
  context: ChatContext;
  agents: AgentInfo[];
  backthink: BackthinkItem[];
  /** Per-session flow graphs, keyed by the graph's own `sessionId` (裁决 16.4.2):
   *  a backthink / sub-session graph no longer overwrites the graph of the
   *  conversation the admin panel is looking at. Bounded by `MAX_GRAPHS`
   *  (store/sessionCache.ts), most recently updated first. */
  flowGraphs: Record<string, FlowGraph>;
  /** Per-session **lane** graphs (切片 FV2): every 路 (leg) of the session's latest
   *  round, oldest first, keyed by the graph's own `turnId`. Real parallel turns
   *  publish one graph per leg; they used to overwrite each other in the single
   *  session bucket. `flowGraphs` keeps its «session → newest graph» meaning (compat),
   *  the Flow page reads this list to draw each leg of the round as its own lane. */
  flowLanes: Record<string, FlowGraph[]>;
  /** Compatibility alias for the admin Flow / Agents pages: the graph of the
   *  *active* session, falling back to the main session's latest one. Kept in sync
   *  by `setFlowGraph` / `setActiveSession` / `resetChat`. */
  flowGraph: FlowGraph | null;
  deviceUpdates: string[];
  /** 唤醒式能力三态（design/14 §5）——当前唤醒状态、上次切换的来源态、最近点亮的工具。 */
  awakenState: AwakenState;
  awakenFrom: AwakenState | null;
  awakenTool: string | null;
  /** 会话编排（design/15 §5）——她是否在忙（不阻塞输入框）。挂起数 `pendingCount`、
   *  在途路数 `activeTurns`、在途回合绑定 `turnBindings` 与六态写入动作由
   *  `store/turnSlice.ts` 提供（AGENTS §10 文件 ≤500 行）。 */
  sessionBusy: boolean;
  pendingCount: number;
  /** 最近一次主代理代审小字（design/11 §4；`approval_notice` 帧，裁决 13.4.3）：
   *  敏感操作经代审后显示一行「xxx 想要干什么，我审查后批准了」。同一时刻只留
   *  最近一条（新的覆盖旧的），且**不持久化** —— 代审是实时事件，刷新后不重放。 */
  approvalNotice: ApprovalNotice | null;

  /** 适配器消息日志实时事件（`channel_persisted` WS 帧，design/16 §4.6 §4）——
   *  隔离会话（群聊/陌生人）落库后秒推，供 IM 适配器日志页增量插入。最新在前，
   *  最多保留最近 ~200 条，按 id 去重。web 聊天无此字段的消费者。 */
  channelEvents: AdminChannelMessage[];
  /** Push one `channel_persisted` message (newest first, dedup by id, capped ~200). */
  pushChannelEvent: (msg: AdminChannelMessage) => void;

  /** The role catalog from `GET /personas` (design/22 §4.6). The role the UI
   *  renders is *derived* from it plus the active session — never stored — so the
   *  picker, the avatar and the mood cannot drift apart. Empty until loaded (or
   *  when the fetch fails), and the UI then keeps its neutral placeholder. */
  personas: PersonaRecord[];
  /** Replace the role catalog (`GET /personas`). */
  setPersonas: (list: PersonaRecord[]) => void;

  setWsStatus: (status: WsStatus) => void;
  setContext: (patch: Partial<ChatContext>) => void;
  /** Record one role's new mood. `personaId` is the role the event named; "" keeps
   *  an unscoped/legacy frame in its own bucket. A null mood drops the bucket. */
  setMood: (mood: string | null, personaId?: string) => void;
  /** Merge role-level moods read from `GET /personas` into the buckets (design/01
   *  §2: mood follows the role, so the catalog is an authoritative seed). Values
   *  already known are overwritten by the server's truth; buckets the seed does
   *  not mention are left alone, so a WS frame that arrived first is never
   *  cleared. Build the seed with `personaMoodSeed`. */
  seedMoods: (seed: Record<string, string>) => void;

  /** Replace the session metadata list. */
  setSessions: (list: SessionMeta[]) => void;
  /** Zero one session's unread badge locally (Android 协同 B3/D4). Called when the
   *  user opens the session and its page has loaded, and on the `session_read`
   *  frame another end broadcasts — a no-op for an already-read or unknown
   *  session, so neither path can invent a row or re-render the list needlessly. */
  clearUnread: (sessionId: string) => void;
  /** Focus a session: stash the list on screen into its own bucket and take the
   *  target session's bucket out (裁决 16.4.5). Never drops a visited session's
   *  loaded page, so switching back shows it immediately while the backend page
   *  is refetched. */
  setActiveSession: (id: string) => void;
  /** Replace the active session's message list with the backend page. A streaming
   *  assistant bubble that has not been persisted yet is carried over (see
   *  `replacePage`). An empty list clears the panel (the fetch-failed path). */
  setMessages: (list: ChatMessage[]) => void;
  /** Prepend an older page (chronological) onto the active session's messages. */
  prependMessages: (list: ChatMessage[]) => void;
  /** Append a message to the bucket its `sessionId` names, or merge by id if it
   *  already exists there. On `message_persisted` it also adopts a same-role local
   *  optimistic bubble (rewriting its uid to the backend id) so a message never
   *  renders twice (design/16 §4.3 dedup). A message without a session id (a local
   *  optimistic bubble) belongs to the active session. */
  appendOrUpdateMessage: (msg: ChatMessage) => void;
  /** Merge one handshake `history` replay page (contract §3.2) into the session the
   *  frame names — up to 50 already-persisted rows, once per (re)connect. Dedup by
   *  id and adopt a local optimistic bubble only on an exact key (`clientMsgId` for
   *  users / `turnId` for assistants), never by the content/time heuristic: this is
   *  the reconnect catch-up path, so a wrong adopt would rewrite a streaming bubble
   *  with an old reply. */
  mergeHistoryPage: (sessionId: string, msgs: ChatMessage[]) => void;
  /** Patch one already-rendered message in place by id — the `message_edited` WS
   *  path. With `sessionId` the patch is confined to that session's bucket; without
   *  one the id is searched across the active list and every bucket (ids are
   *  globally unique). A miss is a no-op, so an edit can never inject a row. */
  updateMessageById: (id: string, patch: Partial<ChatMessage>, sessionId?: string) => void;
  /** Drop one message by id — the `message_deleted` WS path. Same scoping rules as
   *  `updateMessageById`; a miss is a no-op (another session's message, or a
   *  duplicate frame). */
  removeMessageById: (id: string, sessionId?: string) => void;

  addUserMessage: (text: string, mode?: Mode, image?: string, deviceId?: string, clientMsgId?: string) => void;
  beginAssistant: (mode?: Mode, deviceId?: string) => string;
  /** Append a streamed delta to the assistant bubble of `sessionId` (default: the
   *  active session) — a background turn keeps streaming into its own bucket.
   *  给了 `turnId` 就按**二级分桶**写进那一路的气泡（忙时并行下两条回复交错到达，
   *  按列表尾写会把它们糊成一条）；没给则沿用旧的「列表尾气泡」行为。 */
  appendChunk: (delta: string, sessionId?: string, turnId?: string) => void;
  endStream: (sessionId?: string, turnId?: string) => void;
  setThinking: (text: string, sessionId?: string, turnId?: string) => void;
  addTool: (tool: ToolUsed, sessionId?: string, turnId?: string) => void;

  setAgents: (agents: AgentInfo[]) => void;
  upsertAgent: (agent: AgentInfo) => void;
  addBackthink: (item: BackthinkItem) => void;
  revealBackthink: (id: string, conclusion: string) => void;
  markBackthinkPushed: (id: string) => void;
  /** Store one `graph_update` snapshot under the session it belongs to and refresh
   *  the `flowGraph` alias when that session is the one on screen (裁决 16.4.2). */
  setFlowGraph: (graph: FlowGraph, sessionId?: string) => void;
  addDeviceUpdate: (text: string) => void;
  setAwaken: (
    from: AwakenState | undefined,
    to: AwakenState,
    tool?: string | null
  ) => void;
  /** 记录一条主代理代审小字（`approval_notice` WS 帧）。null 表示清除。 */
  setApprovalNotice: (notice: ApprovalNotice | null) => void;

  taskLogs: Record<string, string>;
  appendTaskLog: (taskId: string, text: string) => void;
  clearTaskLog: (taskId: string) => void;

  resetChat: () => void;
}

export const useAppStore = create<AppStore>()(
  persist(
    (set, get) => ({
      sessions: {},
      activeSessionId: "",
      wsStatus: "idle",
      messages: [],
      sessionMessages: {},
      moods: {},
      context: {},
      agents: [
        { id: "main", name: "昔涟", status: "idle", role: "owner" },
      ],
      backthink: [],
      flowGraphs: {},
      flowLanes: {},
      flowGraph: null,
      deviceUpdates: [],
      awakenState: "quiet",
      awakenFrom: null,
      awakenTool: null,
      sessionBusy: false,
      pendingCount: 0,
      approvalNotice: null,
      personas: [],
      taskLogs: {},
      channelEvents: [],

      // 忙时并行切片（在途路数 / 绑定表 / 队列提示 + 六态与回合生命周期动作）。
      ...turnState(),
      ...turnActions(set, get),
      ...deliveryActions(set, get),

      setWsStatus: (status) => set({ wsStatus: status }),
      setPersonas: (list) => set({ personas: list }),
      setContext: (patch) =>
        set((s) => ({ context: { ...s.context, ...patch } })),
      setMood: (mood, personaId = "") =>
        set((s) => {
          const moods = { ...s.moods };
          if (mood) moods[personaId] = mood;
          else delete moods[personaId];
          return { moods };
        }),
      seedMoods: (seed) =>
        set((s) => (Object.keys(seed).length ? { moods: { ...s.moods, ...seed } } : s)),

      setSessions: (list) =>
        set({
          sessions: Object.fromEntries(list.map((s) => [s.id, s])),
        }),
      setActiveSession: (id) =>
        set((s) => {
          if (id === s.activeSessionId) return s;
          // 离开的会话：把屏幕上的列表存回它自己的桶（未落库的流式内容随之保留）。
          // 还没有归属（""）时不建桶，与切换前一致。
          const buckets = s.activeSessionId
            ? withBucket(s.sessionMessages, s.activeSessionId, s.messages)
            : s.sessionMessages;
          const next = { ...buckets };
          const cached = next[id] ?? [];
          delete next[id];
          return {
            activeSessionId: id,
            messages: cached,
            sessionMessages: next,
            flowGraph: currentGraph(s.flowGraphs, id),
          };
        }),
      clearUnread: (sessionId) =>
        set((s) => {
          const meta = s.sessions[sessionId];
          if (!meta?.unreadCount) return s;
          return {
            sessions: { ...s.sessions, [sessionId]: { ...meta, unreadCount: 0 } },
          };
        }),
      setMessages: (list) =>
        set((s) => ({ messages: replacePage(s.messages, list) })),
      prependMessages: (list) =>
        set((s) => {
          const seen = new Set(s.messages.map((m) => m.id));
          const older = list.filter((m) => !seen.has(m.id));
          return older.length ? { messages: [...older, ...s.messages] } : s;
        }),
      appendOrUpdateMessage: (msg) =>
        set((s) =>
          writeSession(s, msg.sessionId ?? s.activeSessionId, (list) =>
            mergeMessage(list, msg)
          )
        ),

      mergeHistoryPage: (sessionId, list) =>
        set((s) => {
          if (!list.length) return s;
          return writeSession(s, sessionId || s.activeSessionId, (cur) =>
            mergeHistory(cur, list)
          );
        }),

      updateMessageById: (id, patch, sessionId) =>
        set((s) =>
          writeSessionById(s, sessionId, (list) => patchMessage(list, id, patch))
        ),

      removeMessageById: (id, sessionId) =>
        set((s) => writeSessionById(s, sessionId, (list) => dropMessage(list, id))),

      addUserMessage: (text, mode, image, deviceId, clientMsgId) =>
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id: uid(),
              role: "user",
              content: text,
              mode,
              image,
              deviceId,
              clientMsgId,
              /** ① 发送中：本地乐观气泡，落库（`message_persisted`）或幂等命中
               *  （`chat_duplicate`）后归一为服务端 id 并转「已送达」。 */
              delivery: "sending",
              createdAt: now(),
            },
          ],
        })),

      beginAssistant: (mode, deviceId) => {
        const id = uid();
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id,
              role: "assistant",
              content: "",
              mode,
              streaming: true,
              tools: [],
              deviceId,
              createdAt: now(),
            },
          ],
        }));
        return id;
      },

      appendChunk: (delta, sessionId, turnId) =>
        set((s) =>
          writeSession(s, sessionId ?? s.activeSessionId, (list) =>
            turnId
              ? patchAssistantForTurn(list, turnId, (m) => ({
                  ...m,
                  content: m.content + delta,
                }))
              : patchAssistantTail(list, (last) => ({
                  ...last,
                  content: last.content + delta,
                }))
          )
        ),

      endStream: (sessionId, turnId) =>
        set((s) =>
          writeSession(s, sessionId ?? s.activeSessionId, (list) =>
            turnId
              ? patchAssistantForTurn(list, turnId, (m) => ({
                  ...m,
                  streaming: false,
                }))
              : patchAssistantTail(list, (last) => ({ ...last, streaming: false }))
          )
        ),

      setThinking: (text, sessionId, turnId) =>
        set((s) =>
          writeSession(s, sessionId ?? s.activeSessionId, (list) =>
            turnId
              ? patchAssistantForTurn(list, turnId, (m) => ({ ...m, thinking: text }))
              : patchAssistantTail(list, (last) => ({ ...last, thinking: text }))
          )
        ),

      addTool: (tool, sessionId, turnId) =>
        set((s) =>
          writeSession(s, sessionId ?? s.activeSessionId, (list) => {
            const add = (last: ChatMessage): ChatMessage => {
              const existing = last.tools ?? [];
              const idx = existing.findIndex((t) => t.name === tool.name);
              const tools =
                idx === -1
                  ? [...existing, tool]
                  : existing.map((t, i) => (i === idx ? { ...t, ...tool } : t));
              return { ...last, tools };
            };
            return turnId
              ? patchAssistantForTurn(list, turnId, add)
              : patchAssistantTail(list, add);
          })
        ),

      setAgents: (agents) => set({ agents }),
      upsertAgent: (agent) =>
        set((s) => {
          const idx = s.agents.findIndex((a) => a.id === agent.id);
          if (idx === -1) return { agents: [...s.agents, agent] };
          const agents = s.agents.slice();
          agents[idx] = { ...agents[idx], ...agent };
          return { agents };
        }),

      addBackthink: (item) =>
        set((s) => ({ backthink: [item, ...s.backthink] })),
      revealBackthink: (id, conclusion) =>
        set((s) => ({
          backthink: s.backthink.map((b) =>
            b.id === id ? { ...b, conclusion, shouldNotify: true } : b
          ),
        })),
      markBackthinkPushed: (id) =>
        set((s) => ({
          backthink: s.backthink.map((b) =>
            b.id === id ? { ...b, pushed: true } : b
          ),
        })),

      setFlowGraph: (graph, sessionId) =>
        set((s) => {
          // 归属会话：图自己带的 sessionId（后端 RunGraph 建图时写入）优先，其次帧壳的，
          // 最后落到当前会话。每张图只进它自己的桶，互不覆盖（裁决 16.4.2）；真并行时
          // 同一条会话的每一路再各占一格（切片 FV2，见 sessionCache.writeFlowGraph）。
          const sid = graph.sessionId ?? sessionId ?? s.activeSessionId;
          const next = writeFlowGraph(s.flowGraphs, s.flowLanes, sid, graph);
          return {
            ...next,
            flowGraph: currentGraph(next.flowGraphs, s.activeSessionId),
          };
        }),
      addDeviceUpdate: (text) =>
        set((s) => ({ deviceUpdates: [text, ...s.deviceUpdates].slice(0, 50) })),
      setAwaken: (from, to, tool) =>
        set((s) => ({
          awakenState: to,
          awakenFrom: from ?? s.awakenFrom,
          awakenTool: tool ?? s.awakenTool,
        })),
      setApprovalNotice: (notice) => set({ approvalNotice: notice }),

      appendTaskLog: (taskId, text) =>
        set((s) => ({
          taskLogs: {
            ...s.taskLogs,
            [taskId]: (s.taskLogs[taskId] ?? "") + text,
          },
        })),
      clearTaskLog: (taskId) =>
        set((s) => {
          if (!(taskId in s.taskLogs)) return s;
          const taskLogs = { ...s.taskLogs };
          delete taskLogs[taskId];
          return { taskLogs };
        }),

      pushChannelEvent: (msg) =>
        set((s) => {
          if (s.channelEvents.some((m) => m.id === msg.id)) return s;
          return { channelEvents: [msg, ...s.channelEvents].slice(0, 200) };
        }),

      resetChat: () =>
        set((s) => ({
          messages: [],
          sessionMessages: {},
          moods: {},
          flowGraphs: {},
          flowLanes: {},
          flowGraph: null,
          backthink: s.backthink,
          awakenState: "quiet",
          awakenFrom: null,
          awakenTool: null,
          sessionBusy: false,
          pendingCount: 0,
          activeTurns: 0,
          turnBindings: {},
          queueNotice: null,
          approvalNotice: null,
        })),
    }),
    persistOptions
  )
);