import { describe, it, expect } from 'vitest';
import { createTestState } from '../__test-utils__';
import { computeCapacity, finalizeCompute } from './computeSystem';
import type { InfrastructureState } from '@ai-tycoon/shared';
import { INITIAL_INFRASTRUCTURE, FLOPS_TO_TOKENS_MULTIPLIER, COMPUTE_SNAPSHOT_INTERVAL, MAX_COMPUTE_HISTORY } from '@ai-tycoon/shared';

function createInfrastructure(overrides: Partial<InfrastructureState> = {}): InfrastructureState {
  return { ...INITIAL_INFRASTRUCTURE, ...overrides };
}

describe('computeCapacity', () => {
  it('splits allocation between training and inference', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.7 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 500,
    });

    const result = computeCapacity(state, infra);

    expect(result.trainingAllocation).toBe(0.7);
    expect(result.inferenceAllocation).toBeCloseTo(0.3);
  });

  it('calculates effectiveTrainingFlops with cross-hardware contribution', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.6 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 400,
    });

    const result = computeCapacity(state, infra);

    // effectiveTraining = 1000 * 0.6 + 400 * 0.6 * 0.3 = 600 + 72 = 672
    expect(result.effectiveTrainingFlops).toBeCloseTo(672);
  });

  it('calculates effectiveInferenceFlops with cross-hardware contribution', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.6 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 400,
    });

    const result = computeCapacity(state, infra);

    // inferenceAlloc = 0.4
    // effectiveInference = (400 * 0.4 + 1000 * 0.4 * 0.5) * 1 = (160 + 200) * 1 = 360
    expect(result.effectiveInferenceFlops).toBeCloseTo(360);
  });

  it('applies research bonuses to inference calculation', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.5 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 1000,
    });
    const bonuses = {
      tokensPerFlopBonus: 0.3,
      inferenceEfficiencyBonus: 0.15,
      energyCostReduction: 0,
      pipelineSpeedBonus: 0,
      trainingSpeedBonus: 0,
      dataQualityBonus: 0,
      sdkCoverageBonus: 0,
      globalCapabilityBonus: 0,
      reasoningBonus: 0,
      codingBonus: 0,
      creativeBonus: 0,
      multimodalBonus: 0,
      agentsBonus: 0,
      reputationBonus: 0,
      safetyBonus: 0,
      autoScalingBonus: 0,
    };

    const result = computeCapacity(state, infra, bonuses);

    // inferenceBoost = 1 + 0.3 + 0.15 = 1.45
    // effectiveInference = (1000 * 0.5 + 1000 * 0.5 * 0.5) * 1.45 = (500 + 250) * 1.45 = 1087.5
    expect(result.effectiveInferenceFlops).toBeCloseTo(1087.5);
  });

  it('does not apply research bonuses to training calculation', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.5 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 1000,
    });
    const bonuses = {
      tokensPerFlopBonus: 0.5,
      inferenceEfficiencyBonus: 0.5,
      energyCostReduction: 0,
      pipelineSpeedBonus: 0,
      trainingSpeedBonus: 0,
      dataQualityBonus: 0,
      sdkCoverageBonus: 0,
      globalCapabilityBonus: 0,
      reasoningBonus: 0,
      codingBonus: 0,
      creativeBonus: 0,
      multimodalBonus: 0,
      agentsBonus: 0,
      reputationBonus: 0,
      safetyBonus: 0,
      autoScalingBonus: 0,
    };

    const result = computeCapacity(state, infra, bonuses);

    // effectiveTraining = 1000 * 0.5 + 1000 * 0.5 * 0.3 = 500 + 150 = 650
    // (same with or without bonuses)
    expect(result.effectiveTrainingFlops).toBeCloseTo(650);
  });

  it('calculates tokensPerSecondCapacity from effectiveInferenceFlops', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.0 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 0,
      totalInferenceFlops: 100,
    });

    const result = computeCapacity(state, infra);

    // inferenceAlloc = 1.0
    // effectiveInference = (100 * 1.0 + 0 * 1.0 * 0.5) * 1 = 100
    // tokensPerSecond = 100 * 26 = 2600
    expect(result.tokensPerSecondCapacity).toBe(100 * FLOPS_TO_TOKENS_MULTIPLIER);
  });

  it('reports totalFlops as sum of training and inference', () => {
    const state = createTestState();
    const infra = createInfrastructure({
      totalTrainingFlops: 750,
      totalInferenceFlops: 250,
    });

    const result = computeCapacity(state, infra);

    expect(result.totalFlops).toBe(1000);
  });

  it('passes through totalVramGB from infrastructure', () => {
    const state = createTestState();
    const infra = createInfrastructure({
      totalVramGB: 512,
    });

    const result = computeCapacity(state, infra);

    expect(result.totalVramGB).toBe(512);
  });

  it('returns all zeros when infrastructure has no hardware', () => {
    const state = createTestState({
      compute: { trainingAllocation: 0.5 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 0,
      totalInferenceFlops: 0,
      totalVramGB: 0,
    });

    const result = computeCapacity(state, infra);

    expect(result.totalFlops).toBe(0);
    expect(result.effectiveTrainingFlops).toBe(0);
    expect(result.effectiveInferenceFlops).toBe(0);
    expect(result.tokensPerSecondCapacity).toBe(0);
  });

  it('handles full training allocation (1.0)', () => {
    const state = createTestState({
      compute: { trainingAllocation: 1.0 },
    });
    const infra = createInfrastructure({
      totalTrainingFlops: 1000,
      totalInferenceFlops: 500,
    });

    const result = computeCapacity(state, infra);

    // effectiveTraining = 1000 * 1.0 + 500 * 1.0 * 0.3 = 1000 + 150 = 1150
    expect(result.effectiveTrainingFlops).toBeCloseTo(1150);
    // inferenceAlloc = 0, so effectiveInference = 0
    expect(result.effectiveInferenceFlops).toBe(0);
    expect(result.tokensPerSecondCapacity).toBe(0);
  });
});

describe('finalizeCompute', () => {
  it('calculates inferenceUtilization as demand / capacity', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0 } }),
      createInfrastructure({ totalInferenceFlops: 100 }),
    );
    // tokensPerSecondCapacity = 100 * 26 = 2600

    const result = finalizeCompute(capacity, 1300, [], 1);

    expect(result.inferenceUtilization).toBeCloseTo(0.5);
  });

  it('clamps utilization to 1 when demand exceeds capacity', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0 } }),
      createInfrastructure({ totalInferenceFlops: 100 }),
    );

    const result = finalizeCompute(capacity, 999999, [], 1);

    expect(result.inferenceUtilization).toBe(1);
  });

  it('sets utilization to 1 when capacity is 0 but demand > 0', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 0, totalInferenceFlops: 0 }),
    );

    const result = finalizeCompute(capacity, 100, [], 1);

    expect(result.inferenceUtilization).toBe(1);
  });

  it('sets utilization to 0 when both capacity and demand are 0', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 0, totalInferenceFlops: 0 }),
    );

    const result = finalizeCompute(capacity, 0, [], 1);

    expect(result.inferenceUtilization).toBe(0);
  });

  it('adds a history snapshot when tickCount is a multiple of COMPUTE_SNAPSHOT_INTERVAL', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 500, totalInferenceFlops: 500 }),
    );

    const tick = COMPUTE_SNAPSHOT_INTERVAL; // 60
    const result = finalizeCompute(capacity, 200, [], tick);

    expect(result.computeHistory).toHaveLength(1);
    expect(result.computeHistory[0].tick).toBe(tick);
    expect(result.computeHistory[0].totalFlops).toBe(capacity.totalFlops);
    expect(result.computeHistory[0].tokensPerSecondDemand).toBe(200);
  });

  it('does not add a snapshot when tickCount is not a multiple of COMPUTE_SNAPSHOT_INTERVAL', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 500, totalInferenceFlops: 500 }),
    );

    const result = finalizeCompute(capacity, 200, [], 1);

    expect(result.computeHistory).toHaveLength(0);
  });

  it('preserves existing history entries', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 500, totalInferenceFlops: 500 }),
    );

    const existingHistory = [
      {
        tick: 60,
        totalFlops: 800,
        effectiveTrainingFlops: 400,
        effectiveInferenceFlops: 400,
        inferenceUtilization: 0.5,
        tokensPerSecondCapacity: 10400,
        tokensPerSecondDemand: 5200,
      },
    ];

    const result = finalizeCompute(capacity, 200, existingHistory, COMPUTE_SNAPSHOT_INTERVAL * 2);

    expect(result.computeHistory).toHaveLength(2);
    expect(result.computeHistory[0].tick).toBe(60);
    expect(result.computeHistory[1].tick).toBe(COMPUTE_SNAPSHOT_INTERVAL * 2);
  });

  it('trims history when it exceeds MAX_COMPUTE_HISTORY', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.5 } }),
      createInfrastructure({ totalTrainingFlops: 100, totalInferenceFlops: 100 }),
    );

    // Create a history that is exactly at the limit
    const fullHistory = Array.from({ length: MAX_COMPUTE_HISTORY }, (_, i) => ({
      tick: (i + 1) * COMPUTE_SNAPSHOT_INTERVAL,
      totalFlops: 200,
      effectiveTrainingFlops: 100,
      effectiveInferenceFlops: 100,
      inferenceUtilization: 0.5,
      tokensPerSecondCapacity: 2600,
      tokensPerSecondDemand: 1300,
    }));

    const nextSnapshotTick = (MAX_COMPUTE_HISTORY + 1) * COMPUTE_SNAPSHOT_INTERVAL;
    const result = finalizeCompute(capacity, 1300, fullHistory, nextSnapshotTick);

    expect(result.computeHistory).toHaveLength(MAX_COMPUTE_HISTORY);
    // Oldest entry should have been shifted out
    expect(result.computeHistory[0].tick).toBe(2 * COMPUTE_SNAPSHOT_INTERVAL);
    // Newest entry should be our new snapshot
    expect(result.computeHistory[result.computeHistory.length - 1].tick).toBe(nextSnapshotTick);
  });

  it('carries capacity fields into the returned ComputeState', () => {
    const capacity = computeCapacity(
      createTestState({ compute: { trainingAllocation: 0.3 } }),
      createInfrastructure({ totalTrainingFlops: 200, totalInferenceFlops: 800, totalVramGB: 256 }),
    );

    const result = finalizeCompute(capacity, 500, [], 1);

    expect(result.totalFlops).toBe(capacity.totalFlops);
    expect(result.totalTrainingFlops).toBe(capacity.totalTrainingFlops);
    expect(result.totalInferenceFlops).toBe(capacity.totalInferenceFlops);
    expect(result.totalVramGB).toBe(capacity.totalVramGB);
    expect(result.effectiveTrainingFlops).toBe(capacity.effectiveTrainingFlops);
    expect(result.effectiveInferenceFlops).toBe(capacity.effectiveInferenceFlops);
    expect(result.tokensPerSecondCapacity).toBe(capacity.tokensPerSecondCapacity);
    expect(result.tokensPerSecondDemand).toBe(500);
  });
});