import { create } from "zustand";
import type { OtherParameters, Store } from "./types.js";
import { immer } from "zustand/middleware/immer";

export const defaultSystemPrompt = `You are a helpful assistant that answers questions based on the provided context. If you don't know the answer, just say that you don't know, don't try to make up an answer.`;
export const defaultParameters = {
  temperature: 0.5,
  max_tokens: 100,
} as OtherParameters;

export const useStore = create<Store>()(
  immer((set, get) => ({
    selectedConversationId: "",
    conversations: [],
    messages: [],
    message: "",
    systemPrompt: defaultSystemPrompt,
    parameters: defaultParameters,
    facts: [],
    loading: false,
    setConversationId: (conversationId) =>
      set((stateDraft) => {
        stateDraft.selectedConversationId = conversationId;
      }),
    setConversationTitle: (conversationTitle) =>
      set((stateDraft) => {
        const conversation = stateDraft.conversations.find(
          (c) => c.id === stateDraft.selectedConversationId,
        );
        if (conversation) {
          conversation.title = conversationTitle;
        }
      }),
    setConversations: (conversations) =>
      set((stateDraft) => {
        stateDraft.conversations = conversations;
      }),
    addConversation: (conversation) =>
      set((stateDraft) => {
        stateDraft.conversations.push(conversation);
      }),
    removeConversation: (conversationId) =>
      set((stateDraft) => {
        stateDraft.conversations.splice(
          stateDraft.conversations.findIndex(
            (conversation) => conversation.id === conversationId,
          ),
          1,
        );
      }),
    setMessages: (messages) =>
      set((stateDraft) => {
        //@ts-ignore
        stateDraft.messages = messages;
      }),
    setMessage: (message) =>
      set((stateDraft) => {
        stateDraft.message = message;
      }),
    setSystemPrompt: (systemPrompt) =>
      set((stateDraft) => {
        stateDraft.systemPrompt = systemPrompt;
      }),
    setParameters: (parameters) =>
      set((stateDraft) => {
        stateDraft.parameters = parameters;
      }),
    setFacts: (facts) =>
      set((stateDraft) => {
        stateDraft.facts = facts;
      }),
    removeFact: (factId) =>
      set((stateDraft) => {
        stateDraft.facts.splice(
          stateDraft.facts.findIndex((fact) => fact.id === factId),
          1,
        );
      }),
    setLoading: (loading) =>
      set((stateDraft) => {
        stateDraft.loading = loading;
      }),
  })),
);