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

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

export interface AppStore {
  sessionId: string;
  wsStatus: WsStatus;
  messages: ChatMessage[];
  mood: string | null;
  context: ChatContext;
  agents: AgentInfo[];
  backthink: BackthinkItem[];
  flowGraph: FlowGraph | null;
  deviceUpdates: string[];

  setWsStatus: (status: WsStatus) => void;
  setSessionId: (id: string) => void;
  setContext: (patch: Partial<ChatContext>) => void;
  setMood: (mood: string | null) => void;

  addUserMessage: (text: string, mode?: Mode) => void;
  beginAssistant: (mode?: Mode) => 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;

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

  resetChat: () => void;
}

export const useAppStore = create<AppStore>()(
  persist(
    (set, get) => ({
      sessionId: "web",
      wsStatus: "idle",
      messages: [],
      mood: null,
      context: {},
      agents: [
        { id: "main", name: "昔涟", status: "idle", role: "owner" },
      ],
      backthink: [],
      flowGraph: null,
      deviceUpdates: [],
      taskLogs: {},

      setWsStatus: (status) => set({ wsStatus: status }),
      setSessionId: (id) => set({ sessionId: id }),
      setContext: (patch) =>
        set((s) => ({ context: { ...s.context, ...patch } })),
      setMood: (mood) => set({ mood }),

      addUserMessage: (text, mode) =>
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id: uid(),
              role: "user",
              content: text,
              mode,
              createdAt: now(),
            },
          ],
        })),

      beginAssistant: (mode) => {
        const id = uid();
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id,
              role: "assistant",
              content: "",
              mode,
              streaming: true,
              tools: [],
              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;
          messages[messages.length - 1] = {
            ...last,
            tools: [...(last.tools ?? []), tool],
          };
          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) })),

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

      resetChat: () =>
        set((s) => ({
          messages: [],
          mood: null,
          flowGraph: null,
          backthink: s.backthink,
        })),
    }),
    {
      name: "cyrene.app",
      version: 1,
      partialize: (s) => ({
        sessionId: s.sessionId,
        messages: s.messages,
        mood: s.mood,
        context: s.context,
        backthink: s.backthink,
      }),
    }
  )
);

/** 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;
}