Zustand Scoped Store：Zustand + ReactContext#

import { createContext, useContext, useRef } from 'react'
import { create } from 'zustand'
import { useStoreWithEqualityFn } from 'zustand/traditional'
import { shallow } from 'zustand/shallow'

// 1) store state + action
interface CounterState {
  count: number
  step: number
  setStep: (n: number) => void
  increment: () => void
  reset: () => void
}

// 2) store 工厂：每次调用可生成独立实例（关键）
const createCounterStore = (initialCount = 0, initialStep = 1) =>
  create<CounterState>()((set, get) => ({
    count: initialCount,
    step: initialStep,
    setStep: (n) => set({ step: Math.max(1, n) }),
    increment: () => set({ count: get().count + get().step }),
    reset: () => set({ count: initialCount }),
  }))

// 3) Context 只负责“传 store 实例”
const CounterStoreContext = createContext<
  ReturnType<typeof createCounterStore> | null
>(null)

interface CounterProviderProps {
  initialCount?: number
  initialStep?: number
  children: React.ReactNode
}

export function CounterProvider({
  initialCount = 0,
  initialStep = 1,
  children,
}: CounterProviderProps) {
  // 用 ref 保证 Provider 生命周期内 store 实例稳定
  const storeRef = useRef<ReturnType<typeof createCounterStore> | null>(null)

  if (!storeRef.current) {
    storeRef.current = createCounterStore(initialCount, initialStep)
  }

  return (
    <CounterStoreContext.Provider value={storeRef.current}>
      {children}
    </CounterStoreContext.Provider>
  )
}

// 4) 统一消费 hook：selector + shallow 做精准订阅
export function useCounterStore<T>(selector: (s: CounterState) => T): T {
  const store = useContext(CounterStoreContext)
  if (!store) {
    throw new Error('useCounterStore must be used within CounterProvider')
  }
  return useStoreWithEqualityFn(store, selector, shallow)
}

// 5) 业务组件：仅订阅自己需要的字段
function CounterPanel() {
  const { count, step, setStep, increment, reset } = useCounterStore((s) => ({
    count: s.count,
    step: s.step,
    setStep: s.setStep,
    increment: s.increment,
    reset: s.reset,
  }))

  return (
    <div style={{ border: '1px solid #ddd', padding: 12, marginBottom: 12 }}>
      <p>count: {count}</p>
      <p>step: {step}</p>
      <button onClick={increment}>+step</button>
      <button onClick={reset} style={{ marginLeft: 8 }}>
        reset
      </button>
      <button onClick={() => setStep(step + 1)} style={{ marginLeft: 8 }}>
        step + 1
      </button>
    </div>
  )
}

// 6) 同页多个 Provider：状态互不干扰（scoped store 的价值）
export default function App() {
  return (
    <>
      <h3>Counter A</h3>
      <CounterProvider initialCount={0} initialStep={1}>
        <CounterPanel />
      </CounterProvider>

      <h3>Counter B</h3>
      <CounterProvider initialCount={100} initialStep={10}>
        <CounterPanel />
      </CounterProvider>
    </>
  )
}