import { query } from '@/modules';
import { message } from 'antd';
import { produce } from 'immer';
import { create } from 'zustand';
import { AiHistoryStore, createAiHistoryStore } from './ai-history';
type ResData = {
  created_at: string;
  done?: boolean;
  done_reason?: string;
  eval_count?: number;
  eval_duration?: number;
  load_duration?: number;
  message?: { role?: string; content?: string }[];
  model?: string;
  prompt_eval_count?: number;
  prompt_eval_duration?: number;
  total_duration?: number;
};
type MessageData = {
  chatId: string;
  chatPromptId: string;
  role: string;
  root: boolean;
  show: boolean;
  data: {
    message: string;
  };
};
export type AiStore = {
  /** 打开侧边栏 */
  open: boolean;
  setOpen: (open: boolean) => void;
  showEdit: boolean;
  setShowEdit: (showEdit: boolean) => void;
  loading: boolean;
  setLoading: (loading: boolean) => void;
  type?: string;
  /** 当前app执行的key */
  key: string;
  setKey: (key: string) => void;
  setType?: (type: string) => void;
  formData: any;
  setFormData: (data: any) => void;
  // Context Data
  data: any;
  setData: (data: any) => void;
  title: string; // AI Module的标题
  setTitle: (title: string) => void;
  messages: { role: string; content: string }[];
  setMessage: (message: { role: string; content: string }[]) => void;
  getPrompt: (key?: string) => any;
  // 设置预设数据，只有是root的情况下才成立
  presetData?: any;
  setPresetData: (data: any) => void;
  /** 是否是第一条消息 */
  root: boolean;
  setRoot: (root: boolean) => void;
  onMessage: (data: MessageData) => void;
} & AiHistoryStore;

export const useAiStore = create<AiStore>((set, get, store) => {
  return {
    open: false,
    setOpen: (open) => set({ open, root: true }),
    loading: false,
    setLoading: (loading) => set({ loading }),
    key: '',
    setKey: (key) => {
      const { getPrompt } = get();
      if (key) {
        getPrompt(key);
      }
      set({ key });
    },
    formData: {},
    setFormData: (data) => set({ formData: data }),
    title: '',
    setTitle: (title) => set({ title }),
    data: {},
    setData: (data) => set({ data }),
    messages: [],
    setMessage: (messages) => set({ messages }),
    getPrompt: async (key) => {
      const state = get();
      const res = await query.post({
        path: 'chat-prompt',
        key: 'getByKey',
        data: {
          key: key || state.key,
        },
      });
      if (res.code === 200) {
        const { prompt } = res.data;
        const { presetData } = prompt;
        console.log(res.data);
        state.setPresetData(presetData);
        set({ root: true });
      }
    },
    presetData: {},
    setPresetData: (data) => {
      set({ presetData: data });
      const { inputs = [] } = data?.data || {};
      const formData = {
        inputs: inputs.map((input) => {
          return {
            key: input.key,
            value: input.value,
            type: 'string',
          };
        }),
        messages: [],
      };
      set({ formData });
    },
    root: false,
    setRoot: (root) => set({ root }),
    onMessage: (data) => {
      const state = get();
      if (state.root) set({ root: false });
      const { message } = data.data;
      const role = data.role;
      set(
        produce((state) => {
          state.messages.push({ role, content: message });
        }),
      );
    },
    showEdit: false,
    setShowEdit: (showEdit) => set({ showEdit }),
    ...createAiHistoryStore(set, get, store),
  };
});
