AIHostingTycoon/apps/web/src/store/index.ts

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import type {
  GameState, Era, GameSpeed, GameSettings,
  EconomyState, InfrastructureState, ComputeState,
  ResearchState, ModelsState, MarketState,
  CompetitorState, TalentState, DataState,
  ReputationState, EventState, AchievementState,
  DataCenter, GpuType, GpuInventory, TrainingJob,
  ActiveResearch, EventConsequence, OwnedDataset,
} from '@ai-tycoon/shared';
import type { FundingRoundType, OverloadPolicy, TuningPreset, ModelTuning } from '@ai-tycoon/shared';
import {
  INITIAL_SETTINGS, SAVE_VERSION,
  INITIAL_ECONOMY, INITIAL_INFRASTRUCTURE, INITIAL_COMPUTE,
  INITIAL_RESEARCH, INITIAL_MODELS, INITIAL_MARKET,
  INITIAL_COMPETITORS, INITIAL_TALENT, INITIAL_DATA,
  INITIAL_REPUTATION, INITIAL_EVENTS, INITIAL_ACHIEVEMENTS,
  GPU_CONFIGS,
  FUNDING_ROUNDS,
  OPEN_SOURCE_REPUTATION_BOOST,
} from '@ai-tycoon/shared';
import { INITIAL_RIVALS } from '@ai-tycoon/game-engine';

export type ActivePage = 'dashboard' | 'infrastructure' | 'research' | 'models'
  | 'market' | 'talent' | 'data' | 'competitors' | 'finance' | 'achievements' | 'settings';

interface UIState {
  activePage: ActivePage;
  notifications: GameNotification[];
}

export interface GameNotification {
  id: string;
  title: string;
  message: string;
  type: 'info' | 'success' | 'warning' | 'danger';
  tick: number;
  read: boolean;
}

interface Actions {
  setActivePage: (page: ActivePage) => void;
  addNotification: (n: Omit<GameNotification, 'id' | 'read'>) => void;
  dismissNotification: (id: string) => void;
  startNewGame: (companyName: string) => void;
  setGameSpeed: (speed: GameSpeed) => void;
  togglePause: () => void;
  setTrainingAllocation: (ratio: number) => void;
  buyGpu: (dataCenterId: string, gpuType: GpuType, count: number) => void;
  buildDataCenter: (name: string, location: DataCenter['location']) => void;
  startTraining: (job: Omit<TrainingJob, 'progressTicks'>) => void;
  deployModel: (modelId: string) => void;
  setProductPricing: (productLineId: string, field: string, value: number) => void;
  toggleProductLine: (productLineId: string) => void;
  startResearch: (research: ActiveResearch) => void;
  resolveEvent: (instanceId: string, choiceIndex: number) => void;
  hireDepartment: (departmentId: string, count: number) => void;
  purchaseDataset: (dataset: OwnedDataset, cost: number) => void;
  raiseFunding: (roundType: FundingRoundType) => void;
  openSourceModel: (modelId: string) => void;
  setOverloadPolicy: (policy: Partial<OverloadPolicy>) => void;
  setModelTuning: (modelId: string, tuning: Partial<ModelTuning>) => void;
  acquireCompetitor: (competitorId: string) => void;
  updateState: (partial: Partial<GameState>) => void;
}

type Store = GameState & UIState & Actions;

const initialGameState: GameState = {
  meta: {
    saveVersion: SAVE_VERSION,
    companyName: '',
    currentEra: 'startup',
    tickCount: 0,
    lastTickTimestamp: Date.now(),
    gameSpeed: 1,
    isPaused: true,
    createdAt: Date.now(),
    totalPlayTime: 0,
    settings: INITIAL_SETTINGS,
  },
  economy: INITIAL_ECONOMY,
  infrastructure: INITIAL_INFRASTRUCTURE,
  compute: INITIAL_COMPUTE,
  research: INITIAL_RESEARCH,
  models: INITIAL_MODELS,
  market: INITIAL_MARKET,
  competitors: INITIAL_COMPETITORS,
  talent: INITIAL_TALENT,
  data: INITIAL_DATA,
  reputation: INITIAL_REPUTATION,
  events: INITIAL_EVENTS,
  achievements: INITIAL_ACHIEVEMENTS,
};

export const useGameStore = create<Store>()(
  persist(
    (set, get) => ({
      ...initialGameState,
      activePage: 'dashboard' as ActivePage,
      notifications: [],

      setActivePage: (page) => set({ activePage: page }),

      addNotification: (n) => set((s) => ({
        notifications: [
          { ...n, id: crypto.randomUUID(), read: false },
          ...s.notifications.slice(0, 49),
        ],
      })),

      dismissNotification: (id) => set((s) => ({
        notifications: s.notifications.map(n =>
          n.id === id ? { ...n, read: true } : n,
        ),
      })),

      startNewGame: (companyName) => set({
        ...initialGameState,
        meta: {
          ...initialGameState.meta,
          companyName,
          isPaused: false,
          createdAt: Date.now(),
          lastTickTimestamp: Date.now(),
        },
        competitors: {
          rivals: INITIAL_RIVALS,
          industryBenchmark: 0,
        },
        activePage: 'dashboard',
        notifications: [],
      }),

      setGameSpeed: (speed) => set((s) => ({
        meta: { ...s.meta, gameSpeed: speed },
      })),

      togglePause: () => set((s) => ({
        meta: { ...s.meta, isPaused: !s.meta.isPaused },
      })),

      setTrainingAllocation: (ratio) => set((s) => ({
        compute: { ...s.compute, trainingAllocation: ratio, inferenceAllocation: 1 - ratio },
      })),

      buyGpu: (dataCenterId, gpuType, count) => set((s) => {
        const price = s.infrastructure.gpuMarketPrices[gpuType] * count;
        if (s.economy.money < price) return s;

        const dataCenters = s.infrastructure.dataCenters.map(dc => {
          if (dc.id !== dataCenterId) return dc;
          const existingIdx = dc.gpus.findIndex(g => g.type === gpuType);
          let gpus: GpuInventory[];
          if (existingIdx >= 0) {
            gpus = dc.gpus.map((g, i) =>
              i === existingIdx
                ? { ...g, count: g.count + count, healthyCount: g.healthyCount + count }
                : g,
            );
          } else {
            gpus = [...dc.gpus, { type: gpuType, count, healthyCount: count, failedCount: 0 }];
          }
          return { ...dc, gpus };
        });

        return {
          economy: { ...s.economy, money: s.economy.money - price },
          infrastructure: { ...s.infrastructure, dataCenters },
        };
      }),

      buildDataCenter: (name, location) => set((s) => {
        const buildCost = 10_000;
        if (s.economy.money < buildCost) return s;

        const dc: DataCenter = {
          id: crypto.randomUUID(),
          name,
          location,
          gpus: [],
          maxCapacity: 100,
          coolingLevel: 0.5,
          redundancyLevel: 0.3,
          currentUptime: 1,
          energyCostPerTick: 0,
          maintenanceCostPerTick: 0,
        };

        return {
          economy: { ...s.economy, money: s.economy.money - buildCost },
          infrastructure: {
            ...s.infrastructure,
            dataCenters: [...s.infrastructure.dataCenters, dc],
          },
        };
      }),

      startTraining: (job) => set((s) => ({
        models: {
          ...s.models,
          activeTraining: { ...job, progressTicks: 0 },
        },
      })),

      deployModel: (modelId) => set((s) => ({
        models: {
          ...s.models,
          trainedModels: s.models.trainedModels.map(m =>
            m.id === modelId ? { ...m, isDeployed: true } : m,
          ),
          productLines: s.models.productLines.map(pl => ({
            ...pl, modelId, isActive: true,
          })),
        },
      })),

      setProductPricing: (productLineId, field, value) => set((s) => ({
        models: {
          ...s.models,
          productLines: s.models.productLines.map(pl =>
            pl.id === productLineId
              ? { ...pl, pricing: { ...pl.pricing, [field]: value } }
              : pl,
          ),
        },
      })),

      toggleProductLine: (productLineId) => set((s) => ({
        models: {
          ...s.models,
          productLines: s.models.productLines.map(pl =>
            pl.id === productLineId ? { ...pl, isActive: !pl.isActive } : pl,
          ),
        },
      })),

      startResearch: (research) => set((s) => {
        if (s.research.activeResearch) return s;
        return {
          research: { ...s.research, activeResearch: research },
        };
      }),

      resolveEvent: (instanceId, choiceIndex) => set((s) => {
        const event = s.events.activeEvents.find(e => e.instanceId === instanceId);
        if (!event) return s;

        const choice = event.choices[choiceIndex];
        if (!choice) return s;

        let money = s.economy.money;
        let reputation = { ...s.reputation };
        const consequences = choice.consequences;

        for (const c of consequences) {
          switch (c.type) {
            case 'money': money += c.value; break;
            case 'reputation': reputation = { ...reputation, score: Math.min(100, Math.max(0, reputation.score + c.value)), publicPerception: Math.min(100, Math.max(0, reputation.publicPerception + c.value)) }; break;
          }
        }

        return {
          economy: { ...s.economy, money: Math.max(0, money) },
          reputation,
          events: {
            ...s.events,
            activeEvents: s.events.activeEvents.filter(e => e.instanceId !== instanceId),
            eventHistory: [
              ...s.events.eventHistory,
              {
                eventId: event.eventId,
                instanceId,
                title: event.title,
                category: event.category,
                tick: s.meta.tickCount,
                chosenOptionIndex: choiceIndex,
              },
            ],
          },
        };
      }),

      hireDepartment: (departmentId, count) => set((s) => {
        const costPerHire = 2000;
        const totalCost = costPerHire * count;
        if (s.economy.money < totalCost) return s;

        return {
          economy: { ...s.economy, money: s.economy.money - totalCost },
          talent: {
            ...s.talent,
            departments: {
              ...s.talent.departments,
              [departmentId]: {
                ...s.talent.departments[departmentId as keyof typeof s.talent.departments],
                headcount: s.talent.departments[departmentId as keyof typeof s.talent.departments].headcount + count,
              },
            },
          },
        };
      }),

      purchaseDataset: (dataset, cost) => set((s) => {
        if (s.economy.money < cost) return s;
        return {
          economy: { ...s.economy, money: s.economy.money - cost },
          data: {
            ...s.data,
            ownedDatasets: [...s.data.ownedDatasets, dataset],
            totalTrainingTokens: s.data.totalTrainingTokens + dataset.sizeTokens,
          },
        };
      }),

      raiseFunding: (roundType) => set((s) => {
        const config = FUNDING_ROUNDS[roundType];
        if (!config) return s;
        const amount = config.amount;
        const dilution = config.dilution;
        return {
          economy: {
            ...s.economy,
            money: s.economy.money + amount,
            funding: {
              ...s.economy.funding,
              totalRaised: s.economy.funding.totalRaised + amount,
              founderEquity: s.economy.funding.founderEquity * (1 - dilution),
              completedRounds: [
                ...s.economy.funding.completedRounds,
                { type: roundType, amount, dilution, completedAtTick: s.meta.tickCount },
              ],
              isPublic: roundType === 'ipo',
            },
          },
        };
      }),

      openSourceModel: (modelId) => set((s) => {
        if (s.market.openSourcedModels.includes(modelId)) return s;
        return {
          market: {
            ...s.market,
            openSourcedModels: [...s.market.openSourcedModels, modelId],
          },
          reputation: {
            ...s.reputation,
            score: Math.min(100, s.reputation.score + OPEN_SOURCE_REPUTATION_BOOST),
            publicPerception: Math.min(100, s.reputation.publicPerception + OPEN_SOURCE_REPUTATION_BOOST),
          },
        };
      }),

      setOverloadPolicy: (policy) => set((s) => ({
        market: {
          ...s.market,
          overloadPolicy: { ...s.market.overloadPolicy, ...policy },
        },
      })),

      setModelTuning: (modelId, tuning) => set((s) => ({
        models: {
          ...s.models,
          trainedModels: s.models.trainedModels.map(m =>
            m.id === modelId ? { ...m, tuning: { ...m.tuning, ...tuning } } : m,
          ),
        },
      })),

      acquireCompetitor: (competitorId) => set((s) => {
        const rival = s.competitors.rivals.find(r => r.id === competitorId);
        if (!rival || rival.status === 'acquired') return s;
        const cost = rival.estimatedRevenue * 500 + rival.estimatedCapability * 100_000;
        if (s.economy.money < cost) return s;
        return {
          economy: { ...s.economy, money: s.economy.money - cost },
          competitors: {
            ...s.competitors,
            rivals: s.competitors.rivals.map(r =>
              r.id === competitorId ? { ...r, status: 'acquired' as const } : r,
            ),
          },
          talent: {
            ...s.talent,
            departments: {
              ...s.talent.departments,
              research: { ...s.talent.departments.research, headcount: s.talent.departments.research.headcount + 5 },
              engineering: { ...s.talent.departments.engineering, headcount: s.talent.departments.engineering.headcount + 3 },
            },
          },
        };
      }),

      updateState: (partial) => set((s) => {
        const newState: Partial<Store> = {};
        for (const key of Object.keys(partial) as (keyof GameState)[]) {
          const value = partial[key];
          const current = s[key];
          if (typeof value === 'object' && value !== null && !Array.isArray(value) && typeof current === 'object' && current !== null) {
            (newState as Record<string, unknown>)[key] = { ...current, ...value };
          } else {
            (newState as Record<string, unknown>)[key] = value;
          }
        }
        return newState;
      }),
    }),
    {
      name: 'ai-tycoon-save',
      partialize: (state) => {
        const { activePage, notifications, ...rest } = state;
        return rest;
      },
    },
  ),
);