import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';
import type { Company } from '@shared/types/entities.types';

interface CompanyState {
  // State
  currentCompany: Company | null;
  companies: Company[];
  isLoading: boolean;

  // Actions
  setCurrentCompany: (company: Company | null) => void;
  setCompanies: (companies: Company[]) => void;
  updateCompany: (data: Partial<Company>) => void;
  clearCompany: () => void;
}

export const useCompanyStore = create<CompanyState>()(
  persist(
    (set, get) => ({
      // Initial state
      currentCompany: null,
      companies: [],
      isLoading: false,

      // Actions
      setCurrentCompany: (company) => {
        set({ currentCompany: company });
      },

      setCompanies: (companies) => {
        set({ companies });
        // If no current company and companies exist, set first one
        if (!get().currentCompany && companies.length > 0) {
          set({ currentCompany: companies[0] });
        }
      },

      updateCompany: (data) => {
        const { currentCompany, companies } = get();
        if (currentCompany) {
          const updated = { ...currentCompany, ...data };
          set({
            currentCompany: updated,
            companies: companies.map((c) =>
              c.id === updated.id ? updated : c
            ),
          });
        }
      },

      clearCompany: () => {
        set({
          currentCompany: null,
          companies: [],
        });
      },
    }),
    {
      name: 'company-storage',
      storage: createJSONStorage(() => localStorage),
      partialize: (state) => ({
        currentCompany: state.currentCompany,
      }),
    }
  )
);
