import { create } from 'zustand';
import { immer } from 'zustand/middleware/immer';
import type { LlmProvidersConfig, LlmProvider, LlmModel } from '@/types';

export type ModalState = 
  | { type: 'closed' }
  | { type: 'editProvider'; providerId: string }
  | { type: 'addProvider' }
  | { type: 'editModel'; providerId: string; modelIndex: number }
  | { type: 'addModel'; providerId: string }
  | { type: 'discoverModels'; providerId: string };

type LlmConfigState = {
  providers: LlmProvidersConfig;
  modal: ModalState;
  setInitialProviders: (providers: LlmProvidersConfig) => void;
  openModal: (modalState: ModalState) => void;
  closeModal: () => void;

  // Provider actions
  updateProvider: (providerId: string, provider: LlmProvider) => void;
  addProvider: (providerId: string, provider: LlmProvider) => void;
  deleteProvider: (providerId: string) => void;

  // Model actions
  updateModel: (providerId: string, modelIndex: number, model: LlmModel) => void;
  addModel: (providerId: string, model: LlmModel) => void;
  deleteModel: (providerId: string, modelIndex: number) => void;
};

export const useLlmConfigStore = create<LlmConfigState>()(
  immer((set) => ({
    providers: {},
    modal: { type: 'closed' },

    setInitialProviders: (providers) => set({ providers }),
    openModal: (modalState) => set({ modal: modalState }),
    closeModal: () => set({ modal: { type: 'closed' } }),

    updateProvider: (providerId, provider) => {
      set((state) => {
        state.providers[providerId] = provider;
      });
    },
    addProvider: (providerId, provider) => {
      set((state) => {
        state.providers[providerId] = provider;
      });
    },
    deleteProvider: (providerId) => {
      set((state) => {
        delete state.providers[providerId];
      });
    },
    updateModel: (providerId, modelIndex, model) => {
      set((state) => {
        state.providers[providerId].models[modelIndex] = model;
      });
    },
    addModel: (providerId, model) => {
      set((state) => {
        state.providers[providerId].models.push(model);
      });
    },
    deleteModel: (providerId, modelIndex) => {
      set((state) => {
        state.providers[providerId].models.splice(modelIndex, 1);
      });
    },
  }))
);
