import { vec2 } from 'gl-matrix';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

import { appConfig } from '../config';
import { AgentGenerationPipeline } from '../pipelines/agents/agent-generation/agent-generation-pipeline';
import { settings } from '../settings';
import { AgentPopulation } from './agent-population';

vi.hoisted(() => {
  Object.defineProperty(globalThis, 'localStorage', {
    configurable: true,
    value: {
      getItem: vi.fn(() => null),
      setItem: vi.fn(),
    },
  });
});

const originalBrushSize = settings.brushSize;
const originalSelectedColorIndex = settings.selectedColorIndex;
const originalSpawnPerPixel = settings.spawnPerPixel;

const createPopulation = () => {
  const pipeline = {
    maxAgentCount: 10_000_000,
    writeAgents: vi.fn(),
    resizeAgents: vi.fn(),
    compactAgents: vi.fn(),
  } as unknown as AgentGenerationPipeline;

  return new AgentPopulation(pipeline);
};

const setPopulationActiveCount = (population: AgentPopulation, activeCount: number) => {
  Object.assign(population as unknown as Record<string, number>, {
    activeCount,
  });
};

const setPopulationAdaptiveCap = (population: AgentPopulation, adaptiveCap: number) => {
  Object.assign(population as unknown as Record<string, number>, {
    adaptiveCap,
  });
};

const getPopulationAdaptiveCap = (population: AgentPopulation): number =>
  (population as unknown as { adaptiveCap: number }).adaptiveCap;

describe('AgentPopulation adaptive budget', () => {
  beforeEach(() => {
    settings.brushSize = 1;
    settings.selectedColorIndex = 0;
    settings.spawnPerPixel = 1;
  });

  afterEach(() => {
    settings.brushSize = originalBrushSize;
    settings.selectedColorIndex = originalSelectedColorIndex;
    settings.spawnPerPixel = originalSpawnPerPixel;
  });

  it('expands beyond the 1M start cap only when new agents arrive under healthy FPS', () => {
    const population = createPopulation();
    setPopulationActiveCount(population, 1_000_000);

    population.growBudget(1 / 60, 60);
    population.spawnStrokeAgents(vec2.fromValues(0, 0), vec2.fromValues(16, 0));

    expect(getPopulationAdaptiveCap(population)).toBeGreaterThan(
      appConfig.simulation.budget.adaptiveCapInitial
    );
    expect(population.activeAgentCount).toBeGreaterThan(
      appConfig.simulation.budget.adaptiveCapInitial
    );
    expect(getPopulationAdaptiveCap(population)).toBeLessThanOrEqual(
      appConfig.simulation.budget.adaptiveCapMax
    );
  });

  it('does not grow the cap above the adaptive max agent count', () => {
    const population = createPopulation();
    const maxAgentCount = appConfig.simulation.budget.adaptiveCapMax;
    setPopulationAdaptiveCap(population, maxAgentCount - 1);
    setPopulationActiveCount(population, maxAgentCount - 1);

    population.growBudget(1 / 60, 60);
    population.spawnStrokeAgents(vec2.fromValues(0, 0), vec2.fromValues(16, 0));

    expect(getPopulationAdaptiveCap(population)).toBe(maxAgentCount);
    expect(population.activeAgentCount).toBe(maxAgentCount);
  });

  it('clamps a stale cap before adding agents', () => {
    const population = createPopulation();
    const maxAgentCount = appConfig.simulation.budget.adaptiveCapMax;
    setPopulationAdaptiveCap(population, maxAgentCount + 1_000);
    setPopulationActiveCount(population, maxAgentCount);

    population.spawnStrokeAgents(vec2.fromValues(0, 0), vec2.fromValues(16, 0));

    expect(getPopulationAdaptiveCap(population)).toBe(maxAgentCount);
    expect(population.activeAgentCount).toBe(maxAgentCount);
  });

  it('decreases the cap and active count slowly when FPS falls below the threshold', () => {
    const population = createPopulation();
    setPopulationActiveCount(population, 1_000_000);

    population.growBudget(10, 50);

    expect(getPopulationAdaptiveCap(population)).toBe(
      appConfig.simulation.budget.adaptiveCapMin
    );
    expect(population.activeAgentCount).toBe(appConfig.simulation.budget.adaptiveCapMin);
  });
});