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, AchievementState,
  Cluster, Campus, DataCenter, DCTier, RackSkuId, TrainingJob,
  ActiveResearch, OwnedDataset, LocationId,
  DeploymentCohort, PipelineStage,
  NetworkHealthState,
} 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_ACHIEVEMENTS,
  DC_TIER_CONFIGS, RACK_SKU_CONFIGS,
  PIPELINE_ORDER_BASE_TICKS, DC_UPGRADE_COST_FRACTION, DC_UPGRADE_INCREMENT,
  CLUSTER_COST_CONFIG, CAMPUS_TIER_COSTS, FIRST_CAMPUS_BUILD_TICKS,
  COHORT_SCALE_FACTOR,
  FUNDING_ROUNDS,
  OPEN_SOURCE_REPUTATION_BOOST,
  LOCATION_CONFIGS,
  networkSlotsRequired, maxComputeRacks,
  uuid,
} 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' | 'leaderboard' | 'settings';

export type InfraNavLevel = 'clusters' | 'cluster' | 'campus' | 'datacenter';

export interface InfraNav {
  level: InfraNavLevel;
  clusterId?: string;
  campusId?: string;
  datacenterId?: string;
}

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

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

function emptyNetworkHealth(): NetworkHealthState {
  return { tier1Required: 0, tier1Healthy: 0, tier2Required: 0, tier2Healthy: 0, tier3Required: 0, tier3Healthy: 0, racksDisconnected: 0 };
}

interface Actions {
  setActivePage: (page: ActivePage) => void;
  setInfraNav: (nav: InfraNav) => void;
  addNotification: (n: Omit<GameNotification, 'id' | 'read'>) => void;
  dismissNotification: (id: string) => void;
  markAllNotificationsRead: () => void;
  startNewGame: (companyName: string) => void;
  setGameSpeed: (speed: GameSpeed) => void;
  togglePause: () => void;
  setTrainingAllocation: (ratio: number) => void;
  buildCluster: (name: string, locationId: LocationId) => void;
  buildCampus: (name: string, clusterId: string, dcTier: DCTier) => void;
  buildDataCenter: (name: string, campusId: string) => void;
  deployRacks: (dataCenterId: string, skuId: RackSkuId, quantity: number) => void;
  fillDCToCapacity: (dataCenterId: string, skuId: RackSkuId) => void;
  addDCsToCampus: (campusId: string, count: number) => void;
  retrofitDC: (dataCenterId: string, newSkuId: RackSkuId) => void;
  cancelRetrofit: (dataCenterId: string) => void;
  upgradeDataCenter: (dataCenterId: string, upgrade: 'cooling' | 'redundancy') => 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;
  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,
  achievements: INITIAL_ACHIEVEMENTS,
};

// --- Helper: find entities in the hierarchy ---

function findCluster(infra: InfrastructureState, clusterId: string): Cluster | undefined {
  return infra.clusters.find(c => c.id === clusterId);
}

function findCampusInCluster(cluster: Cluster, campusId: string): Campus | undefined {
  return cluster.campuses.find(c => c.id === campusId);
}

function findCampus(infra: InfrastructureState, campusId: string): { cluster: Cluster; campus: Campus } | undefined {
  for (const cluster of infra.clusters) {
    const campus = cluster.campuses.find(c => c.id === campusId);
    if (campus) return { cluster, campus };
  }
  return undefined;
}

function findDC(infra: InfrastructureState, dcId: string): { cluster: Cluster; campus: Campus; dc: DataCenter } | undefined {
  for (const cluster of infra.clusters) {
    for (const campus of cluster.campuses) {
      const dc = campus.dataCenters.find(d => d.id === dcId);
      if (dc) return { cluster, campus, dc };
    }
  }
  return undefined;
}

function updateDCInInfra(infra: InfrastructureState, dcId: string, updater: (dc: DataCenter) => DataCenter): InfrastructureState {
  return {
    ...infra,
    clusters: infra.clusters.map(cluster => ({
      ...cluster,
      campuses: cluster.campuses.map(campus => ({
        ...campus,
        dataCenters: campus.dataCenters.map(dc =>
          dc.id === dcId ? updater(dc) : dc,
        ),
      })),
    })),
  };
}

function updateCampusInInfra(infra: InfrastructureState, campusId: string, updater: (campus: Campus) => Campus): InfrastructureState {
  return {
    ...infra,
    clusters: infra.clusters.map(cluster => ({
      ...cluster,
      campuses: cluster.campuses.map(campus =>
        campus.id === campusId ? updater(campus) : campus,
      ),
    })),
  };
}

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

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

      setInfraNav: (nav) => set({ infraNav: nav }),

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

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

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

      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: [],
        infraNav: { level: 'clusters' },
      }),

      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 },
      })),

      // --- Infrastructure: Cluster ---

      buildCluster: (name, locationId) => set((s) => {
        const loc = LOCATION_CONFIGS[locationId];
        const eraOrder: Era[] = ['startup', 'scaleup', 'bigtech', 'agi'];
        if (eraOrder.indexOf(s.meta.currentEra) < eraOrder.indexOf(loc.availableAt)) return s;

        const existingInRegion = s.infrastructure.clusters.find(c => c.locationId === locationId);
        if (existingInRegion) return s;

        const isFirst = s.infrastructure.clusters.length === 0;
        const cost = isFirst ? 0 : CLUSTER_COST_CONFIG.baseCost;
        if (s.economy.money < cost) return s;

        const cluster: Cluster = {
          id: uuid(),
          name,
          locationId,
          campuses: [],
          status: isFirst ? 'operational' : 'constructing',
          constructionProgress: isFirst ? 0 : 0,
          constructionTotal: isFirst ? 0 : CLUSTER_COST_CONFIG.buildTimeTicks,
        };

        return {
          economy: { ...s.economy, money: s.economy.money - cost },
          infrastructure: {
            ...s.infrastructure,
            clusters: [...s.infrastructure.clusters, cluster],
          },
        };
      }),

      // --- Infrastructure: Campus ---

      buildCampus: (name, clusterId, dcTier) => set((s) => {
        const cluster = findCluster(s.infrastructure, clusterId);
        if (!cluster || cluster.status !== 'operational') return s;

        const tierConfig = DC_TIER_CONFIGS[dcTier];
        const eraOrder: Era[] = ['startup', 'scaleup', 'bigtech', 'agi'];
        if (eraOrder.indexOf(s.meta.currentEra) < eraOrder.indexOf(tierConfig.requiredEra)) return s;
        if (tierConfig.requiredResearch && !s.research.completedResearch.includes(tierConfig.requiredResearch)) return s;

        const campusCost = CAMPUS_TIER_COSTS[dcTier];
        const isFirstCampus = s.infrastructure.clusters.every(c => c.campuses.length === 0);
        const cost = isFirstCampus ? 0 : campusCost.baseCost;
        if (s.economy.money < cost) return s;

        const buildTime = isFirstCampus ? FIRST_CAMPUS_BUILD_TICKS : campusCost.buildTimeTicks;

        const campus: Campus = {
          id: uuid(),
          name,
          clusterId,
          dcTier,
          dataCenters: [],
          status: 'constructing',
          constructionProgress: 0,
          constructionTotal: buildTime,
        };

        return {
          economy: { ...s.economy, money: s.economy.money - cost },
          infrastructure: {
            ...s.infrastructure,
            clusters: s.infrastructure.clusters.map(c =>
              c.id === clusterId
                ? { ...c, campuses: [...c.campuses, campus] }
                : c,
            ),
          },
        };
      }),

      // --- Infrastructure: Data Center ---

      buildDataCenter: (name, campusId) => set((s) => {
        const found = findCampus(s.infrastructure, campusId);
        if (!found || found.campus.status !== 'operational') return s;

        const tierConfig = DC_TIER_CONFIGS[found.campus.dcTier];
        if (s.economy.money < tierConfig.baseCost) return s;

        const isFirstDC = s.infrastructure.clusters.every(cl =>
          cl.campuses.every(ca => ca.dataCenters.length === 0),
        );
        const buildTime = isFirstDC ? tierConfig.firstBuildTimeTicks : tierConfig.buildTimeTicks;

        const dc: DataCenter = {
          id: uuid(),
          name,
          campusId,
          tier: found.campus.dcTier,
          status: 'constructing',
          constructionProgress: 0,
          constructionTotal: buildTime,
          rackSkuId: null,
          computeRacksOnline: 0,
          computeRacksFailed: 0,
          networkHealth: emptyNetworkHealth(),
          deploymentCohorts: [],
          retrofitState: null,
          coolingLevel: 0,
          redundancyLevel: 0,
          effectiveComputeRacks: 0,
          usedSlots: 0,
          usedPowerKW: 0,
          energyCostPerTick: 0,
          maintenanceCostPerTick: 0,
          currentUptime: 1,
        };

        return {
          economy: { ...s.economy, money: s.economy.money - tierConfig.baseCost },
          infrastructure: updateCampusInInfra(s.infrastructure, campusId, (campus) => ({
            ...campus,
            dataCenters: [...campus.dataCenters, dc],
          })),
        };
      }),

      // --- Infrastructure: Deploy Racks ---

      deployRacks: (dataCenterId, skuId, quantity) => set((s) => {
        if (quantity <= 0) return s;

        const found = findDC(s.infrastructure, dataCenterId);
        if (!found || found.dc.status !== 'operational') return s;

        const dc = found.dc;
        if (dc.rackSkuId !== null && dc.rackSkuId !== skuId) return s;

        const sku = RACK_SKU_CONFIGS[skuId];
        const eraOrder: Era[] = ['startup', 'scaleup', 'bigtech', 'agi'];
        if (eraOrder.indexOf(s.meta.currentEra) < eraOrder.indexOf(sku.era)) return s;
        if (sku.requiredResearch && !s.research.completedResearch.includes(sku.requiredResearch)) return s;

        const tierConfig = DC_TIER_CONFIGS[dc.tier];
        const maxCompute = maxComputeRacks(tierConfig.rackSlots);
        const existingCompute = dc.computeRacksOnline + dc.computeRacksFailed
          + dc.deploymentCohorts.filter(c => c.stage !== 'decommission').reduce((sum, c) => sum + c.count, 0);
        const available = maxCompute - existingCompute;
        const actualQty = Math.min(quantity, available);
        if (actualQty <= 0) return s;

        const totalNetSlots = networkSlotsRequired(existingCompute + actualQty);
        const totalSlotsNeeded = existingCompute + actualQty + totalNetSlots;
        if (totalSlotsNeeded > tierConfig.rackSlots) return s;

        const powerNeeded = (existingCompute + actualQty) * sku.powerDrawKW;
        if (powerNeeded > tierConfig.powerBudgetKW) return s;

        const totalCost = sku.baseCost * actualQty;
        if (s.economy.money < totalCost) return s;

        const baseTicks = PIPELINE_ORDER_BASE_TICKS;
        const scaledTicks = Math.ceil(baseTicks * (1 + COHORT_SCALE_FACTOR * actualQty));

        const cohort: DeploymentCohort = {
          id: uuid(),
          count: actualQty,
          skuId,
          stage: 'ordered',
          stageProgress: 0,
          stageTotal: scaledTicks,
          repairCount: 0,
        };

        return {
          economy: { ...s.economy, money: s.economy.money - totalCost },
          infrastructure: updateDCInInfra(s.infrastructure, dataCenterId, (d) => ({
            ...d,
            rackSkuId: skuId,
            deploymentCohorts: [...d.deploymentCohorts, cohort],
          })),
        };
      }),

      fillDCToCapacity: (dataCenterId, skuId) => {
        const s = get();
        const found = findDC(s.infrastructure, dataCenterId);
        if (!found || found.dc.status !== 'operational') return;

        const dc = found.dc;
        const tierConfig = DC_TIER_CONFIGS[dc.tier];
        const maxCompute = maxComputeRacks(tierConfig.rackSlots);
        const existingCompute = dc.computeRacksOnline + dc.computeRacksFailed
          + dc.deploymentCohorts.filter(c => c.stage !== 'decommission').reduce((sum, c) => sum + c.count, 0);
        const available = maxCompute - existingCompute;
        if (available <= 0) return;

        const sku = RACK_SKU_CONFIGS[skuId];
        const affordableQty = Math.floor(s.economy.money / sku.baseCost);
        const powerLimit = Math.floor((tierConfig.powerBudgetKW - existingCompute * sku.powerDrawKW) / sku.powerDrawKW);
        const qty = Math.min(available, affordableQty, powerLimit);
        if (qty <= 0) return;

        get().deployRacks(dataCenterId, skuId, qty);
      },

      addDCsToCampus: (campusId, count) => set((s) => {
        if (count <= 0) return s;

        const found = findCampus(s.infrastructure, campusId);
        if (!found || found.campus.status !== 'operational') return s;

        const tierConfig = DC_TIER_CONFIGS[found.campus.dcTier];
        const totalCost = tierConfig.baseCost * count;
        if (s.economy.money < totalCost) return s;

        const newDCs: DataCenter[] = [];
        for (let i = 0; i < count; i++) {
          newDCs.push({
            id: uuid(),
            name: `${found.campus.name}-DC-${found.campus.dataCenters.length + i + 1}`,
            campusId,
            tier: found.campus.dcTier,
            status: 'constructing',
            constructionProgress: 0,
            constructionTotal: tierConfig.buildTimeTicks,
            rackSkuId: null,
            computeRacksOnline: 0,
            computeRacksFailed: 0,
            networkHealth: emptyNetworkHealth(),
            deploymentCohorts: [],
            retrofitState: null,
            coolingLevel: 0,
            redundancyLevel: 0,
            effectiveComputeRacks: 0,
            usedSlots: 0,
            usedPowerKW: 0,
            energyCostPerTick: 0,
            maintenanceCostPerTick: 0,
            currentUptime: 1,
          });
        }

        return {
          economy: { ...s.economy, money: s.economy.money - totalCost },
          infrastructure: updateCampusInInfra(s.infrastructure, campusId, (campus) => ({
            ...campus,
            dataCenters: [...campus.dataCenters, ...newDCs],
          })),
        };
      }),

      // --- Infrastructure: Retrofit ---

      retrofitDC: (dataCenterId, newSkuId) => set((s) => {
        const found = findDC(s.infrastructure, dataCenterId);
        if (!found || found.dc.status !== 'operational') return s;

        const dc = found.dc;
        if (!dc.rackSkuId || dc.rackSkuId === newSkuId) return s;

        const sku = RACK_SKU_CONFIGS[newSkuId];
        const eraOrder: Era[] = ['startup', 'scaleup', 'bigtech', 'agi'];
        if (eraOrder.indexOf(s.meta.currentEra) < eraOrder.indexOf(sku.era)) return s;
        if (sku.requiredResearch && !s.research.completedResearch.includes(sku.requiredResearch)) return s;

        const totalRacksToRetrofit = dc.computeRacksOnline + dc.computeRacksFailed;
        if (totalRacksToRetrofit <= 0) return s;

        const oldSku = RACK_SKU_CONFIGS[dc.rackSkuId];
        const decommTicks = Math.ceil(oldSku.pipelineTimeTicks.installation * (1 + COHORT_SCALE_FACTOR * totalRacksToRetrofit));

        return {
          infrastructure: updateDCInInfra(s.infrastructure, dataCenterId, (d) => ({
            ...d,
            status: 'retrofitting' as const,
            deploymentCohorts: [],
            retrofitState: {
              fromSkuId: d.rackSkuId!,
              toSkuId: newSkuId,
              phase: 'decommissioning' as const,
              progress: 0,
              total: decommTicks,
              racksRemaining: totalRacksToRetrofit,
            },
          })),
        };
      }),

      cancelRetrofit: (dataCenterId) => set((s) => {
        const found = findDC(s.infrastructure, dataCenterId);
        if (!found || found.dc.status !== 'retrofitting') return s;

        return {
          infrastructure: updateDCInInfra(s.infrastructure, dataCenterId, (d) => ({
            ...d,
            status: 'operational' as const,
            retrofitState: null,
          })),
        };
      }),

      // --- Infrastructure: Upgrades ---

      upgradeDataCenter: (dataCenterId, upgrade) => set((s) => {
        const found = findDC(s.infrastructure, dataCenterId);
        if (!found || found.dc.status !== 'operational') return s;

        const dc = found.dc;
        const tierConfig = DC_TIER_CONFIGS[dc.tier];
        const cost = tierConfig.baseCost * DC_UPGRADE_COST_FRACTION;
        if (s.economy.money < cost) return s;

        const currentLevel = upgrade === 'cooling' ? dc.coolingLevel : dc.redundancyLevel;
        if (currentLevel >= 1.0) return s;

        return {
          economy: { ...s.economy, money: s.economy.money - cost },
          infrastructure: updateDCInInfra(s.infrastructure, dataCenterId, (d) => ({
            ...d,
            [upgrade === 'cooling' ? 'coolingLevel' : 'redundancyLevel']:
              Math.min(1.0, currentLevel + DC_UPGRADE_INCREMENT),
          })),
        };
      }),

      // --- Non-infrastructure actions (unchanged) ---

      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 },
        };
      }),

      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',
      version: SAVE_VERSION,
      partialize: (state) => {
        const { activePage, notifications, infraNav, ...rest } = state;
        return rest;
      },
      migrate: (_persisted, version) => {
        if (version < SAVE_VERSION) {
          return {
            ...initialGameState,
            activePage: 'dashboard' as const,
            notifications: [{
              id: uuid(),
              title: 'Save Reset',
              message: 'Your save was reset due to a major infrastructure redesign — Hypercluster scale! Build clusters, campuses, and data centers.',
              type: 'info' as const,
              tick: 0,
              read: false,
            }],
            infraNav: { level: 'clusters' },
          } as unknown as Store;
        }
        return _persisted as Store;
      },
    },
  ),
);