import { create } from 'zustand'

const useStore = create((set) => ({
  count: 1,
  // 通過 set 方法更新狀態值，set 支持傳入函數和狀態對象值
  inc: () => set((state) => ({ count: state.count + 1 })),
}))

function Counter() {
  const count = useStore((state) => state.count)
  const inc = useStore((state) => state.inc)
  return (
    <div>
      <p>{`Count: ${count}`}</p>
      <button onClick={inc}>+1</button>
    </div>
  )
}

<div>
  <p>Count: <span id="value"></span></p>
  <button id="btn">+1</button>
</div>

import { createStore } from 'zustand@4.5.4/vanilla'

const store = createStore((set) => ({
  count: 1,
  inc: () => set((state) => ({ count: state.count + 1 })),
}))

const { getState, setState, subscribe, getInitialState } = store

window.onload = () => {
  const value = document.querySelector('#value')
  value.innerHTML = getInitialState().count // 設置 store 中 count 值
  // 使用 subscribe 訂閱狀態變化，并更新數值
  subscribe((state) => {
    value.innerHTML = state.count
  })

  const btn = document.querySelector('#btn')
  btn.onclick = () => {
    // 觸發更新
    getState().inc()
  }
}

const create = (createState) => {
  let state
  let initialState
  const listeners = new Set()

  const setState = (partial, replace) => {
    // 判斷是否為函數，為函數就調用，并傳入當前狀態值
    const nextState = typeof partial === 'function'
      ? partial(state)
      : state
    
    // 對比狀態值是否有變化
    if (!Object.is(nextState, state)) {
      const previousState = state
      // 如果是替換整個狀態值，或者狀態值為基礎值或 null，則直接賦值，不然使用 Object.aasign 合并狀態值
      state = replace ?? (typeof nextState !== 'object' || nextState === null)
        ? nextState
        : Object.assign({}, state, nextState)
      
      // 觸發訂閱函數
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState = () => state
  
  const getInitialState = () => initialState

  const subscribe = (listener) => {
    listeners.add(listener)
    // 返回一個取消訂閱的方法
    return () => {
      listeners.delete(listener)
    }
  }

  // 清空訂閱
  const destory = () => listeners.clear()

  const api = {
    setState,
    getState,
    getInitialState,
    subscribe,
    destory
  }

  // 調用 createState，createState 參數為 set、get 和 api 對象，函數返回狀態初始值
  initialState = (state = createState(setState, getState, api))

  return api
}

export default create

const snapshot = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)

import createImpl from './vanilla'
import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'

const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports

const create = (createState) => {
  // 使用普通版本 zustand 創建一個支持發布訂閱的數據源
  const api = createImpl(createState)

  // zustand 版本 Hook，參數為狀態值選擇器和判斷狀態是否變化函數
  const useBearStore = (selector, equiltyFn?) => 
    // 和 useSyncExternalStore 類似，不過支持傳入 selector，獲取部分數據
    useSyncExternalStoreWithSelector(
      api.subcribe,
      api.getState,
      api.getInitialState,
      selector,
      equiltyFn,
    )

  // 把 api 合并到 Hook 對象上
  Object.assign(useBearStore, api)

  return useBearStore
}

export default create

// 相比于 useSyncExternalStore ，多了 selector 和 isEqual 參數
function useSyncExternalStoreWithSelector(
  subscribe,
  getSnapshot,
  getServerSnapshot,
  selector,
  isEqual?,
) {
  const [getSelection, getServerSelection] = useMemo(() => {
    let memoizedSnapshot; // 緩存的整個狀態值
    let memoizedSelection: Selection; // 緩存的使用 selector 選中的部分狀態值
  
    const memoizedSelector = (nextSnapshot: Snapshot) => {
      const prevSnapshot = memoizedSnapshot
      const prevSelection = memoizedSelection

      // 如果整體狀態值相等，直接返回緩存的 selector 選中的狀態值。
      if (is(prevSnapshot, nextSnapshot)) {
        return prevSelection;
      }

      // 使用 selector 函數獲取最新的狀態值
      const nextSelection = selector(nextSnapshot);

      // 有傳入判斷狀態是否相等函數，相等的話就返回上次的 selector 選中值。
      if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
        // 記錄最新的整體狀態值
        memoizedSnapshot = nextSnapshot;
        return prevSelection;
      }

      // 記錄最新一次的更新值
      memoizedSnapshot = nextSnapshot;
      memoizedSelection = nextSelection;
      // 返回 selector 函數獲取最新的狀態值
      return nextSelection;
    };
      
    const getSnapshotWithSelector = () => memoizedSelector(getSnapshot());
    return [getSnapshotWithSelector, () => {}];
  }, [getSnapshot, getServerSnapshot, selector, isEqual]);
  
  // 調用 useSyncExternalStore 方法，第二參數不是整體獲取整個狀態值，而是 selector 的狀態值
  const value = useSyncExternalStore(
    subscribe,
    getSelection,
    getServerSelection,
  );
  return value
}