import { create } from 'zustand';
import { createJSONStorage, persist } from 'zustand/middleware';

export const useRepoStore = create(
  persist(
    (set) => ({
  repoId: null,
  repoUrl: '',
  ingestionData: null,
  isIngesting: false,
  ingestionProgress: [],
  architectureData: null,
  graphData: null,
  metricsData: null,
  isAnalyzing: false,
  chatHistory: [],
  isAsking: false,
  impactData: null,
  impactHistory: [],
  isAnalyzingImpact: false,
  docsContent: null,
  docsByType: {},
  isGeneratingDocs: false,

  setRepoId: (id) => set({ repoId: id }),
  setRepoUrl: (url) => set({ repoUrl: url }),
  setIngestionData: (data) => set({ ingestionData: data }),
  setIngesting: (value) => set({ isIngesting: value }),
  addProgress: (message) => set((state) => ({ ingestionProgress: [...state.ingestionProgress, message] })),
  resetProgress: () => set({ ingestionProgress: [] }),
  setArchitectureData: (data) => set({ architectureData: data }),
  setGraphData: (data) => set({ graphData: data }),
  setMetricsData: (data) => set({ metricsData: data }),
  setAnalyzing: (value) => set({ isAnalyzing: value }),
  addChatMessage: (message) => set((state) => ({ chatHistory: [...state.chatHistory, message] })),
  clearChat: () => set({ chatHistory: [] }),
  setAsking: (value) => set({ isAsking: value }),
  setImpactData: (data) =>
    set((state) => ({
      impactData: data,
      impactHistory: data ? [data, ...state.impactHistory].slice(0, 6) : state.impactHistory,
    })),
  setAnalyzingImpact: (value) => set({ isAnalyzingImpact: value }),
  setDocsContent: (data) =>
    set((state) => ({
      docsContent: data,
      docsByType: data?.doc_type ? { ...state.docsByType, [data.doc_type]: data } : state.docsByType,
    })),
  setGeneratingDocs: (value) => set({ isGeneratingDocs: value }),
  resetAll: () =>
    set({
      repoId: null,
      repoUrl: '',
      ingestionData: null,
      architectureData: null,
      graphData: null,
      metricsData: null,
      chatHistory: [],
      impactData: null,
      impactHistory: [],
      docsContent: null,
      docsByType: {},
    }),
    }),
    {
      name: 'codeatlas-repo-state',
      storage: createJSONStorage(() => window.sessionStorage),
      partialize: (state) => ({
        repoId: state.repoId,
        repoUrl: state.repoUrl,
        ingestionData: state.ingestionData,
        architectureData: state.architectureData,
        graphData: state.graphData,
        metricsData: state.metricsData,
        chatHistory: state.chatHistory,
        impactData: state.impactData,
        impactHistory: state.impactHistory,
        docsContent: state.docsContent,
        docsByType: state.docsByType,
      }),
    },
  ),
);