import type { GameState, ComputeState, InfrastructureState } from '@token-empire/shared';
import { FLOPS_TO_TOKENS_MULTIPLIER, COMPUTE_SNAPSHOT_INTERVAL, MAX_COMPUTE_HISTORY } from '@token-empire/shared';
import type { ResearchBonuses } from './researchBonuses';

export interface CapacityResult {
  totalFlops: number;
  totalTrainingFlops: number;
  totalInferenceFlops: number;
  totalVramGB: number;
  trainingAllocation: number;
  inferenceAllocation: number;
  effectiveTrainingFlops: number;
  effectiveInferenceFlops: number;
  tokensPerSecondCapacity: number;
}

export function computeCapacity(state: GameState, infrastructure: InfrastructureState, researchBonuses?: ResearchBonuses): CapacityResult {
  const { totalTrainingFlops, totalInferenceFlops, totalVramGB } = infrastructure;
  const trainingAllocation = state.compute.trainingAllocation;
  const inferenceAllocation = 1 - trainingAllocation;

  const effectiveTrainingFlops =
    totalTrainingFlops * trainingAllocation +
    totalInferenceFlops * trainingAllocation * 0.3;

  const inferenceBoost = 1 + (researchBonuses?.tokensPerFlopBonus ?? 0) + (researchBonuses?.inferenceEfficiencyBonus ?? 0);
  const effectiveInferenceFlops =
    (totalInferenceFlops * inferenceAllocation +
    totalTrainingFlops * inferenceAllocation * 0.5) * inferenceBoost;

  const tokensPerSecondCapacity = effectiveInferenceFlops * FLOPS_TO_TOKENS_MULTIPLIER;

  return {
    totalFlops: totalTrainingFlops + totalInferenceFlops,
    totalTrainingFlops,
    totalInferenceFlops,
    totalVramGB,
    trainingAllocation,
    inferenceAllocation,
    effectiveTrainingFlops,
    effectiveInferenceFlops,
    tokensPerSecondCapacity,
  };
}

export function finalizeCompute(capacity: CapacityResult, totalTokenDemand: number, prevHistory: ComputeState['computeHistory'], tickCount: number): ComputeState {
  const inferenceUtilization = capacity.tokensPerSecondCapacity > 0
    ? Math.min(1, totalTokenDemand / capacity.tokensPerSecondCapacity)
    : (totalTokenDemand > 0 ? 1 : 0);

  const computeHistory = [...prevHistory];
  if (tickCount % COMPUTE_SNAPSHOT_INTERVAL === 0) {
    computeHistory.push({
      tick: tickCount,
      totalFlops: capacity.totalFlops,
      effectiveTrainingFlops: capacity.effectiveTrainingFlops,
      effectiveInferenceFlops: capacity.effectiveInferenceFlops,
      inferenceUtilization,
      tokensPerSecondCapacity: capacity.tokensPerSecondCapacity,
      tokensPerSecondDemand: totalTokenDemand,
    });
    if (computeHistory.length > MAX_COMPUTE_HISTORY) {
      computeHistory.shift();
    }
  }

  return {
    ...capacity,
    tokensPerSecondDemand: totalTokenDemand,
    inferenceUtilization,
    computeHistory,
  };
}

export function processCompute(state: GameState, infrastructure: InfrastructureState): ComputeState {
  const cap = computeCapacity(state, infrastructure);
  return finalizeCompute(cap, state.compute.tokensPerSecondDemand, state.compute.computeHistory, state.meta.tickCount);
}