import { create } from "zustand";
import { persist } from "zustand/middleware";
import type {
  AdminChannelMessage,
  AgentInfo,
  AwakenState,
  BackthinkItem,
  ChatMessage,
  FlowGraph,
  Mode,
  PersonaRecord,
  SessionMeta,
  ToolUsed,
  WsStatus,
} from "../lib/types";
import { resolveSessionPersona } from "../lib/persona";
import { now, uid } from "../lib/util";

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

export interface AppStore {
  /** 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). Reset on
   *  session switch and refilled from the backend (design/16 §4/§5). */
  messages: ChatMessage[];
  /** 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). */
  moods: Record<string, string>;
  context: ChatContext;
  agents: AgentInfo[];
  backthink: BackthinkItem[];
  flowGraph: FlowGraph | null;
  deviceUpdates: string[];
  /** 唤醒式能力三态（design/14 §5）——当前唤醒状态、上次切换的来源态、最近点亮的工具。 */
  awakenState: AwakenState;
  awakenFrom: AwakenState | null;
  awakenTool: string | null;
  /** 会话编排（design/15 §5）——她是否在忙 + 挂起消息数（不阻塞输入框）。 */
  sessionBusy: boolean;
  pendingCount: number;

  /** 适配器消息日志实时事件（`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;

  /** Replace the session metadata list. */
  setSessions: (list: SessionMeta[]) => void;
  /** Re-bind one session to a role locally — the optimistic write and the rollback
   *  target of `lib/sessionPersona.ts#setSessionPersona`. `""` re-binds it to
   *  "inherit the default role". A no-op for an unknown session or one already
   *  bound to that role. */
  setSessionPersona: (sessionId: string, personaId: string) => 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 and clear its message cache (refilled by the caller). */
  setActiveSession: (id: string) => void;
  /** Replace the active session's message list (initial page on switch). */
  setMessages: (list: ChatMessage[]) => void;
  /** Prepend an older page (chronological) onto the active session's messages. */
  prependMessages: (list: ChatMessage[]) => void;
  /** Append a message, or merge by id if it already exists. 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). */
  appendOrUpdateMessage: (msg: ChatMessage) => void;
  /** Patch one already-rendered message in place by id — the `message_edited` WS
   *  path. A no-op when the id is not in the active session, so an edit broadcast
   *  for another conversation can never inject a row here. */
  updateMessageById: (id: string, patch: Partial<ChatMessage>) => void;
  /** Drop one message by id — the `message_deleted` WS path. A no-op when the id
   *  is absent (another session's message, or a duplicate frame). */
  removeMessageById: (id: string) => void;

  addUserMessage: (text: string, mode?: Mode, image?: string, deviceId?: string) => void;
  beginAssistant: (mode?: Mode, deviceId?: string) => string;
  appendChunk: (delta: string) => void;
  endStream: () => void;
  setThinking: (text: string) => void;
  addTool: (tool: ToolUsed) => void;

  setAgents: (agents: AgentInfo[]) => void;
  upsertAgent: (agent: AgentInfo) => void;
  addBackthink: (item: BackthinkItem) => void;
  revealBackthink: (id: string, conclusion: string) => void;
  markBackthinkPushed: (id: string) => void;
  setFlowGraph: (graph: FlowGraph) => void;
  addDeviceUpdate: (text: string) => void;
  setAwaken: (
    from: AwakenState | undefined,
    to: AwakenState,
    tool?: string | null
  ) => void;
  setSessionBusy: (busy: boolean, pendingCount?: number) => 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: [],
      moods: {},
      context: {},
      agents: [
        { id: "main", name: "昔涟", status: "idle", role: "owner" },
      ],
      backthink: [],
      flowGraph: null,
      deviceUpdates: [],
      awakenState: "quiet",
      awakenFrom: null,
      awakenTool: null,
      sessionBusy: false,
      pendingCount: 0,
      personas: [],
      taskLogs: {},
      channelEvents: [],

      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 };
        }),

      setSessions: (list) =>
        set({
          sessions: Object.fromEntries(list.map((s) => [s.id, s])),
        }),
      setSessionPersona: (sessionId, personaId) =>
        set((s) => {
          const meta = s.sessions[sessionId];
          if (!meta || (meta.personaId ?? "") === personaId) return s;
          return {
            sessions: { ...s.sessions, [sessionId]: { ...meta, personaId } },
          };
        }),
      setActiveSession: (id) => set({ activeSessionId: id, messages: [] }),
      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({ 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 { messages: [...older, ...s.messages] };
        }),
      appendOrUpdateMessage: (msg) =>
        set((s) => {
          // Exact id match → merge. Keeps local streaming/tools/thinking, takes
          // the backend's content/mode/mood as truth.
          const idx = s.messages.findIndex((m) => m.id === msg.id);
          if (idx !== -1) {
            const messages = s.messages.slice();
            messages[idx] = { ...messages[idx], ...msg };
            return { messages };
          }
          // No exact id: adopt a local optimistic bubble (no sessionId) of the
          // same role, so a message the local browser already rendered
          // optimistically is never duplicated when message_persisted arrives
          // (design/16 §4.3 dedup by id). Prefer a content match, else the
          // nearest by createdAt; scan newest-first.
          let best = -1;
          let bestDist = Number.MAX_SAFE_INTEGER;
          for (let i = s.messages.length - 1; i >= 0; i--) {
            const m = s.messages[i];
            if (m.sessionId || m.role !== msg.role || m.id === msg.id) continue;
            const dist = m.content === msg.content ? 0 : Math.abs(m.createdAt - msg.createdAt);
            if (dist < bestDist) {
              bestDist = dist;
              best = i;
            }
          }
          if (best !== -1) {
            const m = s.messages[best];
            const messages = s.messages.slice();
            messages[best] = {
              ...m,
              id: msg.id,
              sessionId: msg.sessionId ?? m.sessionId,
              content: msg.content || m.content,
              mode: msg.mode ?? m.mode,
              mood: msg.mood ?? m.mood,
              image: msg.image ?? m.image,
              platform: msg.platform ?? m.platform,
              deviceId: msg.deviceId ?? m.deviceId,
              personaId: msg.personaId ?? m.personaId,
              createdAt: msg.createdAt ?? m.createdAt,
            };
            return { messages };
          }
          // No local bubble (e.g. a remote browser that never optimistically
          // rendered this message): append.
          return { messages: [...s.messages, msg] };
        }),

      updateMessageById: (id, patch) =>
        set((s) => {
          const idx = s.messages.findIndex((m) => m.id === id);
          if (idx === -1) return s;
          const messages = s.messages.slice();
          messages[idx] = { ...messages[idx], ...patch };
          return { messages };
        }),

      removeMessageById: (id) =>
        set((s) => {
          const messages = s.messages.filter((m) => m.id !== id);
          return messages.length === s.messages.length ? s : { messages };
        }),

      addUserMessage: (text, mode, image, deviceId) =>
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id: uid(),
              role: "user",
              content: text,
              mode,
              image,
              deviceId,
              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) =>
        set((s) => {
          const messages = s.messages.slice();
          const last = messages[messages.length - 1];
          if (!last || last.role !== "assistant") return s;
          messages[messages.length - 1] = { ...last, content: last.content + delta };
          return { messages };
        }),

      endStream: () =>
        set((s) => {
          const messages = s.messages.slice();
          const last = messages[messages.length - 1];
          if (!last || last.role !== "assistant") return s;
          messages[messages.length - 1] = { ...last, streaming: false };
          return { messages };
        }),

      setThinking: (text) =>
        set((s) => {
          const messages = s.messages.slice();
          const last = messages[messages.length - 1];
          if (!last || last.role !== "assistant") return s;
          messages[messages.length - 1] = { ...last, thinking: text };
          return { messages };
        }),

      addTool: (tool) =>
        set((s) => {
          const messages = s.messages.slice();
          const last = messages[messages.length - 1];
          if (!last || last.role !== "assistant") return s;
          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));
          messages[messages.length - 1] = { ...last, tools };
          return { messages };
        }),

      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) => set({ flowGraph: graph }),
      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,
        })),
      setSessionBusy: (busy, pendingCount) =>
        set((s) => ({
          sessionBusy: busy,
          pendingCount: pendingCount ?? s.pendingCount,
        })),

      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: [],
          moods: {},
          flowGraph: null,
          backthink: s.backthink,
          awakenState: "quiet",
          awakenFrom: null,
          awakenTool: null,
          sessionBusy: false,
          pendingCount: 0,
        })),
    }),
    {
      name: "cyrene.app",
      version: 3,
      // v1 persisted the whole chat (`messages` + a single `sessionId`); v2 added
      // one resolved `currentPersona`, which can no longer describe the active role
      // now that the role belongs to a *session*. Both shapes are dropped:
      // `personas` is refilled from `GET /personas` and the role is re-derived from
      // the session list. `moods` is transient by nature (a restart resets a role's
      // mood), so it is not carried across either. A v1 state was never trusted for
      // the session metadata, so those users simply restart on the main session.
      migrate: (persisted, version) => {
        const p = (persisted ?? {}) as Partial<AppStore>;
        const fromV1 = version < 2;
        return {
          activeSessionId: fromV1 ? "" : p.activeSessionId ?? "",
          sessions: fromV1 ? {} : p.sessions ?? {},
          personas: p.personas ?? [],
          moods: {},
          context: p.context ?? {},
          backthink: p.backthink ?? [],
        };
      },
      partialize: (s) => ({
        // messages deliberately excluded: the backend is the message truth
        // (design/16 §1), so we never persist the whole chat into localStorage.
        // We keep the (small) session metadata + the focused session id so a
        // refresh lands on the same conversation, plus the role catalog so the
        // avatar/name/wallpaper resolve before the first fetch returns.
        activeSessionId: s.activeSessionId,
        sessions: s.sessions,
        personas: s.personas,
        moods: s.moods,
        context: s.context,
        backthink: s.backthink,
      }),
    }
  )
);

/** The role the UI renders (design/22 §4.6): the active session's bound role,
 *  else the global default. THE single resolution point — the avatar, name,
 *  wallpaper, mood bucket and message labels all read it, so none of them can
 *  disagree about "who is speaking" (design/22 §5 risk row 1). */
export function selectCurrentPersona(s: AppStore): PersonaRecord | null {
  return resolveSessionPersona(s.activeSessionId, s.sessions, s.personas);
}

/** Subscribe to the active session's role. */
export function useCurrentPersona(): PersonaRecord | null {
  return useAppStore(selectCurrentPersona);
}

/** Subscribe to the active role's mood. `mood` frames are role-scoped and carry
 *  no sessionId, so a role's mood is only shown above its own sessions. */
export function useCurrentMood(): string | null {
  const personaId = useAppStore((s) => selectCurrentPersona(s)?.id ?? "");
  return useAppStore((s) => (personaId ? s.moods[personaId] ?? null : null));
}

/** Convenience selector for the last streaming assistant message id. */
export function lastStreamingMessage(s: AppStore): ChatMessage | undefined {
  const last = s.messages[s.messages.length - 1];
  return last?.role === "assistant" ? last : undefined;
}

/** Does an edit/delete frame scoped to `sessionId` belong to the conversation the
 *  UI is showing? `message_edited` / `message_deleted` carry the owning session so
 *  the rule matches `message_persisted`: only the active session's messages are
 *  rendered, so another conversation's frame must not touch this list.
 *
 *  A frame without a session id is id-scoped (the backend could not resolve the
 *  conversation — see the gateway's optional message lookup). It still applies:
 *  message ids are globally unique, so the by-id update/removal can only ever match
 *  a message of its own session, and an id we do not hold is a no-op either way. */
export function messageEventApplies(
  sessionId: string | undefined,
  activeSessionId: string
): boolean {
  if (!sessionId) return true;
  return sessionId === activeSessionId;
}