import { createContext, useContext, useMemo, type ReactNode } from "react";
import { create } from "zustand";

export type PeekPayload =
  | { kind: "payer"; payerId: string }
  | { kind: "rule"; rule: string };

interface DrillState {
  stack: PeekPayload[];
  openPeek: (p: PeekPayload) => void;
  closeTop: () => void;
  closeAll: () => void;
}

// One zustand store per provider instance (factory) so multiple
// providers (e.g. in tests) don't share state.
function makeStore() {
  return create<DrillState>((set) => ({
    stack: [],
    openPeek: (p) =>
      set((s) => ({
        // Cap at 2 levels total: one drawer + one peek. When called and
        // the stack already has one peek, replace it.
        stack: s.stack.length >= 1 ? [p] : [p],
      })),
    closeTop: () => set((s) => ({ stack: s.stack.slice(0, -1) })),
    closeAll: () => set({ stack: [] }),
  }));
}

type StoreApi = ReturnType<typeof makeStore>;

const Ctx = createContext<StoreApi | null>(null);

export function DrillStackProvider({ children }: { children: ReactNode }) {
  // useMemo so the store instance is stable across renders.
  const store = useMemo(makeStore, []);
  return <Ctx.Provider value={store}>{children}</Ctx.Provider>;
}

export function useDrillStack() {
  const store = useContext(Ctx);
  if (!store) throw new Error("useDrillStack must be used within DrillStackProvider");
  // Subscribe to just `stack` so consumers re-render only on stack
  // changes (not on every state update).
  const stack = store((s) => s.stack);
  return {
    stack,
    openPeek: store.getState().openPeek,
    closeTop: store.getState().closeTop,
    closeAll: store.getState().closeAll,
  };
}