Mana-Loop/src/lib/game/stores/manaStore.ts

// ─── Mana Store ───────────────────────────────────────────────────────────────
// Handles raw mana, elements, meditation, and mana regeneration
//
// NEW MODEL: All conversion is passive through the unified conversion system.
// convertMana, processConvertAction, and craftComposite are removed (no-ops for save compat).

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { ELEMENTS, BASE_UNLOCKED_ELEMENTS } from '../constants';
import type { ElementState } from '../types';
import { ok, okVoid, fail, ErrorCode } from '../utils/result';
import { createSafeStorage } from '../utils/safe-persist';
import type { Result } from '../utils/result';

// ─── Mana State (data only) ─────────────────────────────────────────────────

export interface ManaState {
  rawMana: number;
  meditateTicks: number;
  totalManaGathered: number;
  elements: Record<string, ElementState>;
  /** Per-element net regen rates (from unified conversion system) */
  elementRegen: Record<string, number>;
}

// ─── Mana Actions ────────────────────────────────────────────────────────────

export interface ManaActions {
  setRawMana: (amount: number) => void;
  addRawMana: (amount: number, maxMana: number) => void;
  spendRawMana: (amount: number) => boolean;
  gatherMana: (amount: number, maxMana: number) => void;

  // Meditation
  setMeditateTicks: (ticks: number) => void;
  incrementMeditateTicks: () => void;
  resetMeditateTicks: () => void;

  // Elements
  unlockElement: (element: string, cost: number) => Result<void>;
  addElementMana: (element: string, amount: number, max: number) => void;
  spendElementMana: (element: string, amount: number) => Result<void>;
  setElementMax: (max: number) => void;
  computeElementMaxWithBonuses: (perElementBonuses: Record<string, number>) => void;
  setElementRegen: (regen: Record<string, number>) => void;

  // Reset
  resetMana: (prestigeUpgrades: Record<string, number>) => void;
}

// ─── Combined Mana Store Type ────────────────────────────────────────────────

export type ManaStore = ManaState & ManaActions;

export const useManaStore = create<ManaStore>()(
  persist(
    (set, get) => ({
      rawMana: 10,
      meditateTicks: 0,
      totalManaGathered: 0,
      elements: Object.fromEntries(
        Object.keys(ELEMENTS).map(k => [
          k,
          {
            current: 0,
            max: 10,
            baseMax: 10,
            unlocked: BASE_UNLOCKED_ELEMENTS.includes(k),
          }
        ])
      ) as Record<string, ElementState>,
      elementRegen: {},

      setRawMana: (amount: number) => {
        set({ rawMana: Math.max(0, amount) });
      },

      addRawMana: (amount: number, maxMana: number) => {
        set((state) => ({
          rawMana: Math.min(state.rawMana + amount, maxMana),
          totalManaGathered: state.totalManaGathered + amount,
        }));
      },

      spendRawMana: (amount: number) => {
        const state = get();
        if (state.rawMana < amount) return false;
        set({ rawMana: state.rawMana - amount });
        return true;
      },

      gatherMana: (amount: number, maxMana: number) => {
        set((state) => ({
          rawMana: Math.min(state.rawMana + amount, maxMana),
          totalManaGathered: state.totalManaGathered + amount,
        }));
      },

      setMeditateTicks: (ticks: number) => set({ meditateTicks: ticks }),
      incrementMeditateTicks: () => set((s) => ({ meditateTicks: s.meditateTicks + 1 })),
      resetMeditateTicks: () => set({ meditateTicks: 0 }),

      unlockElement: (element: string, cost: number) => {
        const state = get();
        if (state.elements[element]?.unlocked) return fail(ErrorCode.INVALID_INPUT, `Element ${element} is already unlocked`);
        if (state.rawMana < cost) return fail(ErrorCode.INSUFFICIENT_MANA, `Need ${cost} raw mana, have ${state.rawMana}`);

        // If the element doesn't exist in the store (e.g. from an old save), create it
        const existing = state.elements[element];
        const newElement = existing
          ? { ...existing, unlocked: true }
          : { current: 0, max: 10, baseMax: 10, unlocked: true };

        set({ rawMana: state.rawMana - cost, elements: { ...state.elements, [element]: newElement } });
        return okVoid();
      },

      addElementMana: (element: string, amount: number, max: number) => {
        set((state) => {
          const elem = state.elements[element];
          if (!elem) return state;
          return {
            elements: { ...state.elements, [element]: { ...elem, current: Math.min(elem.current + amount, max) } },
          };
        });
      },

      spendElementMana: (element: string, amount: number) => {
        const state = get();
        const elem = state.elements[element];
        if (!elem) return fail(ErrorCode.INVALID_ELEMENT, `Element ${element} does not exist`);
        if (elem.current < amount) return fail(ErrorCode.INSUFFICIENT_MANA, `Need ${amount} ${element} mana, have ${elem.current}`);

        set({ elements: { ...state.elements, [element]: { ...elem, current: elem.current - amount } } });
        return okVoid();
      },

      setElementMax: (max: number) => {
        set((state) => ({
          elements: Object.fromEntries(Object.entries(state.elements).map(([k, v]) => [k, { ...v, max, baseMax: v.baseMax ?? max }])) as Record<string, ElementState>,
        }));
      },

      computeElementMaxWithBonuses: (perElementBonuses: Record<string, number>) => {
        set((state) => {
          const newElements = { ...state.elements };
          let changed = false;
          for (const [element, bonus] of Object.entries(perElementBonuses)) {
            if (newElements[element] && bonus > 0) {
              const newMax = (newElements[element].baseMax ?? newElements[element].max) + bonus;
              if (newElements[element].max !== newMax) {
                newElements[element] = { ...newElements[element], max: newMax };
                changed = true;
              }
            }
          }
          return changed ? { elements: newElements } : state;
        });
      },

      setElementRegen: (regen: Record<string, number>) => {
        set({ elementRegen: regen });
      },

      resetMana: (prestigeUpgrades: Record<string, number>) => {
        const elementMax = 10 + (prestigeUpgrades.elementalAttune || 0) * 25;
        const startingMana = 10 + (prestigeUpgrades.manaStart || 0) * 10;
        set({ rawMana: startingMana, meditateTicks: 0, totalManaGathered: 0, elements: makeInitialElements(elementMax, prestigeUpgrades), elementRegen: {} });
      },
    }),
    {
      storage: createSafeStorage(),
      name: 'mana-loop-mana',
      version: 2,
      partialize: (state) => ({ rawMana: state.rawMana, meditateTicks: state.meditateTicks, totalManaGathered: state.totalManaGathered, elements: state.elements, elementRegen: state.elementRegen }),
      migrate: (persistedState: any, _version) => {
        if (persistedState && persistedState.elements) {
          for (const k of Object.keys(persistedState.elements)) {
            if (persistedState.elements[k].baseMax === undefined) {
              persistedState.elements[k].baseMax = persistedState.elements[k].max ?? 10;
            }
          }
          // Add any missing elements that exist in ELEMENTS but not in the save
          for (const k of Object.keys(ELEMENTS)) {
            if (!persistedState.elements[k]) {
              const isUnlocked = BASE_UNLOCKED_ELEMENTS.includes(k);
              persistedState.elements[k] = {
                current: isUnlocked ? 0 : 0,
                max: 10,
                baseMax: 10,
                unlocked: isUnlocked,
              };
            }
          }
        }
        return persistedState;
      },
    }
  )
);

// Helper function to create initial elements
export function makeInitialElements(
  elementMax: number,
  prestigeUpgrades: Record<string, number> = {}
): Record<string, ElementState> {
  const elemStart = (prestigeUpgrades.elemStart || 0) * 5;
  const elements: Record<string, ElementState> = {};
  for (const k of Object.keys(ELEMENTS)) {
    const isUnlocked = BASE_UNLOCKED_ELEMENTS.includes(k);
    elements[k] = { current: isUnlocked ? elemStart : 0, max: elementMax, baseMax: elementMax, unlocked: isUnlocked };
  }
  return elements;
}