/**
 * Copyright 2025 Ori Cohen https://github.com/ori88c
 * https://github.com/ori88c/zero-overhead-keyed-promise-lock
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * General note on testing concurrency in unit tests:
 * While ideal tests follow a strict Arrange-Act-Assert structure, rigorously testing
 * concurrency-oriented components often requires validating intermediate states.
 * Incorrect intermediate states can compromise the entire component's correctness,
 * making their verification essential.
 * 
 * As with everything in engineering, this comes at a cost: verbosity. 
 * Given that resilience is the primary goal, this is a small price to pay.
 */

import { ZeroOverheadKeyedLock } from './zero-overhead-keyed-promise-lock';

const sleep = (ms: number) => new Promise(res => setTimeout(res, ms));

interface MockErrorMessageInfo {
  round: number;
  key: string;
}

describe('ZeroOverheadLock tests', () => {
  describe('Happy path tests', () => {
    test('executeExclusive: should return the expected value when succeeds', async () => {
      // Arrange.
      const lock = new ZeroOverheadKeyedLock<number>();
      const key = 'user-9799789';
      const expectedValue = -295;
      const task = async (): Promise<number> => { return expectedValue; };

      // Pre-action validations.
      expect(lock.isActiveKey(key)).toBe(false);
      expect(lock.getCurrentExecution(key)).toBeUndefined();
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);

      // Act.
      const actualValue = await lock.executeExclusive(key, task);
      
      // Post-action validations.
      expect(actualValue).toBe(expectedValue);
      expect(lock.isActiveKey(key)).toBe(false);
      expect(lock.getCurrentExecution(key)).toBeUndefined();
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
    });
  
    test('waitForAllExistingTasksToComplete: should resolve immediately if no task currently executes', async () => {
      const lock = new ZeroOverheadKeyedLock<void>();
      await lock.waitForAllExistingTasksToComplete();
      expect(lock.getCurrentExecution('mock_key')).toBeUndefined();
      expect(lock.activeKeysCount).toBe(0);
    });

    test(
      'executeExclusive: should process only one task at a time per key, and waitForAllExistingTasksToComplete ' +
      'should resolve only after *all* the currently pending and processed tasks are completed', async () => {
      jest.useFakeTimers();
      const lock = new ZeroOverheadKeyedLock<void>();
      const numberOfKeys = 176;
      const keys: readonly string[] = new Array<number>(numberOfKeys)
        .fill(0)
        .map((_, i) => i*20)
        .map(num => `user-${num}`);
      const roundToExecuteExclusivePromises: Promise<void>[][] = [];
      const initialBackpressurePerKey = 15;
      const taskDurationMs = 8000;
      let actualCompletedTasksCount = 0;
      const createTask = async (): Promise<void> => {
        await sleep(taskDurationMs);
        ++actualCompletedTasksCount;
      };
      const validateKeysActivity = (shouldBePresent: boolean): void => {
        for (const key of keys) {
          expect(lock.isActiveKey(key)).toBe(shouldBePresent);

          const currentExecutionPromise = lock.getCurrentExecution(key);
          if (shouldBePresent) {
            expect(currentExecutionPromise).toBeDefined();
          } else {
            expect(currentExecutionPromise).toBeUndefined();
          }
        }
      };

      // Initial validation.
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      validateKeysActivity(false);

      // Create a burst of tasks, inducing backpressure on the lock.
      // Add `initialBackpressurePerKey` pending tasks per key.
      for (let round = 0; round < initialBackpressurePerKey; ++round) {
        const pendingExecuteExclusivePromises: Promise<void>[] = keys.map(
          key => lock.executeExclusive(key, createTask)
        );
        roundToExecuteExclusivePromises.push(pendingExecuteExclusivePromises);

        // Trigger an event-loop iteration.
        await jest.advanceTimersByTimeAsync(0);

        validateKeysActivity(true);
        expect(lock.activeKeysCount).toBe(keys.length);
        expect(lock.activeKeys).toEqual(keys);
      }

      // Currrently, the number of pending tasks is initialBackpressurePerKey * numberOfKeys.
      let allTasksCompleted = false;
      const waitForCompletionOfAllTasksPromise: Promise<void> = (async () => {
        await lock.waitForAllExistingTasksToComplete();
        allTasksCompleted = true;
      })();

      let expectedCompletedTasksCount = 0;
      for (let round = 0; round < initialBackpressurePerKey; ++round) {
        expect(allTasksCompleted).toBe(false);
        expect(lock.activeKeysCount).toBe(keys.length);
        expect(lock.activeKeys).toEqual(keys);
        validateKeysActivity(true);

        // Simulate the completion of one task per key. This occurs because all tasks 
        // have the same duration, and each key has an equal number of remaining tasks.
        await Promise.race([
          waitForCompletionOfAllTasksPromise,
          jest.advanceTimersByTimeAsync(taskDurationMs)
        ]);
        await Promise.all(roundToExecuteExclusivePromises[round]);

        // Each round, we complete one pending task per key.
        expectedCompletedTasksCount += numberOfKeys;
        expect(actualCompletedTasksCount).toBe(expectedCompletedTasksCount);
      }

      await waitForCompletionOfAllTasksPromise;
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      expect(allTasksCompleted).toBe(true);
      validateKeysActivity(false);
      jest.useRealTimers();
    });

    test('getCurrentExecution should return the active task promise during execution', async () => {
      type ResultType = Record<string, number>;
      const mockResult: ResultType = { a: 1, b: 2 };

      let resolveTask: (value?: unknown) => void;
      const lock = new ZeroOverheadKeyedLock<ResultType>();
      const key = 'mock-key';

      const taskPromise = new Promise<ResultType>(res => resolveTask = res);
      const executeExclusivePromise = lock.executeExclusive(key, () => taskPromise);

      const validationRounds = 24;
      let currentExecution: Promise<ResultType>;
      for (let attempt = 0; attempt < validationRounds; ++attempt) {
        expect(lock.isActiveKey(key)).toBe(true);

        // getCurrentExecution should return the same Promise instance as long as
        // the current task is ongoing.
        if (currentExecution) {
          expect(lock.getCurrentExecution(key)).toBe(currentExecution);
        } else {
          currentExecution = lock.getCurrentExecution(key);
          expect(currentExecution).toBeDefined();
        }

        expect(lock.activeKeysCount).toBe(1);
        expect(lock.activeKeys).toEqual([key]);
      }

      resolveTask(mockResult);
      const result = await executeExclusivePromise;
      expect(result).toBe(mockResult);

      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      expect(lock.isActiveKey(key)).toBe(false);
      expect(lock.getCurrentExecution(key)).toBeUndefined();
    });
  });

  describe('Negative path tests', () => {
    test('executeExclusive & getCurrentExecution: should reject when the key is empty or not a string', async () => {
      const lock = new ZeroOverheadKeyedLock<string>();
      const createTask = () => Promise.resolve<string>('mock-result');
      const invalidKeys = [
        '',
        0 as unknown as string,
        4 as unknown as string,
        -53.431 as unknown as string,
        undefined as string,
        null as string,
        true as unknown as string,
        { prop1: 'value1' } as unknown as string
      ];

      for (const key of invalidKeys) {
        await expect(() => lock.executeExclusive(key, createTask)).rejects.toThrow();
        expect(() => lock.getCurrentExecution(key)).toThrow();
      }

      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      await lock.waitForAllExistingTasksToComplete();
    });

    test('executeExclusive: should return the expected error when throws', async () => {
      const lock = new ZeroOverheadKeyedLock<string>();
      const key = 'mock-key';
      const expectedError = new Error('mock error');
      const createTask = async (): Promise<string> => { throw expectedError; };

      expect.assertions(5);
      try {
        await lock.executeExclusive(key, createTask);
      } catch (err) {
        expect(err).toBe(expectedError);

        // The event-driven eviction mechanism should remove the key after the task completes,
        // as long as no other task is pending (i.e., no backpressure).
        expect(lock.isActiveKey(key)).toBe(false);
        expect(lock.getCurrentExecution(key)).toBeUndefined();
      }

      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      await lock.waitForAllExistingTasksToComplete();
    });

    test(
      'executeExclusive: should process only one task at a time per key, and waitForAllExistingTasksToComplete ' +
      'should resolve only after *all* the currently pending and processed tasks are completed, ' +
      'with all tasks rejecting', async () => {
      jest.useFakeTimers();
      const lock = new ZeroOverheadKeyedLock<void>();
      const numberOfKeys = 96;
      const keys: readonly string[] = new Array<number>(numberOfKeys)
        .fill(0)
        .map((_, i) => i*25)
        .map(num => `user-${num}`);
      const roundToExecuteExclusivePromises: Promise<void>[][] = [];
      const initialBackpressurePerKey = 12;
      const taskDurationMs = 6000;
      const createErrorMessage = (info: MockErrorMessageInfo): string =>
        `Task for ${info.key} failed during round ${info.round}`;
      const task = async (info: MockErrorMessageInfo): Promise<void> => {
        await sleep(taskDurationMs);
        throw new Error(createErrorMessage(info));
      };
      const validateKeysActivity = (shouldBePresent: boolean): void => {
        for (const key of keys) {
          expect(lock.isActiveKey(key)).toBe(shouldBePresent);

          const currentExecutionPromise = lock.getCurrentExecution(key);
          if (shouldBePresent) {
            expect(currentExecutionPromise).toBeDefined();
          } else {
            expect(currentExecutionPromise).toBeUndefined();
          }
        }
      };

      // Initial validation.
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      validateKeysActivity(false);

      // Create a burst of tasks, inducing backpressure on the lock.
      // Add `initialBackpressurePerKey` pending tasks per key.
      for (let round = 0; round < initialBackpressurePerKey; ++round) {
        const pendingExecuteExclusivePromises: Promise<void>[] = keys.map(
          key => lock.executeExclusive(key, () => task({ round, key }))
        );
        roundToExecuteExclusivePromises.push(pendingExecuteExclusivePromises);

        // Trigger an event-loop iteration.
        await jest.advanceTimersByTimeAsync(0);

        validateKeysActivity(true);
        expect(lock.activeKeysCount).toBe(keys.length);
        expect(lock.activeKeys).toEqual(keys);
      }

      // Currrently, the number of pending tasks is initialBackpressurePerKey * numberOfKeys.
      let allTasksCompleted = false;
      const waitForCompletionOfAllTasksPromise: Promise<void> = (async () => {
        await lock.waitForAllExistingTasksToComplete();
        allTasksCompleted = true;
      })();

      for (let round = 0; round < initialBackpressurePerKey; ++round) {
        expect(allTasksCompleted).toBe(false);
        expect(lock.activeKeysCount).toBe(keys.length);
        expect(lock.activeKeys).toEqual(keys);
        validateKeysActivity(true);

        // Simulate the completion of one task per key. This occurs because all tasks 
        // have the same duration, and each key has an equal number of remaining tasks.
        await Promise.allSettled([
          ...roundToExecuteExclusivePromises[round],
          jest.advanceTimersByTimeAsync(taskDurationMs)
        ]);

        // Each round, we complete one pending task per key.
        let i = 0;
        for (const key of keys) {
          try {
            await roundToExecuteExclusivePromises[round][i++];
            expect(true).toBe(false); // The flow should not reach this point.
          } catch (err) {
            const expectedMessage = createErrorMessage({ round, key });
            expect(err.message).toEqual(expectedMessage);
          }
        }
      }

      await waitForCompletionOfAllTasksPromise;
      expect(lock.activeKeysCount).toBe(0);
      expect(lock.activeKeys).toEqual([]);
      expect(allTasksCompleted).toBe(true);
      validateKeysActivity(false);
      jest.useRealTimers();
    });
  });
});