import { Config } from "@jest/types";
import TestRunner, { Test, TestEvents, TestRunnerContext, TestWatcher, UnsubscribeFn } from "jest-runner";

export default class CustomTestRunner {
    public supportsEventEmitters = true;

    private parallelRunner: ParallelTestRunner;
    private sequentialRunner: SequentialTestRunner;

    constructor(globalConfig: Config.GlobalConfig, context: TestRunnerContext) {
        this.parallelRunner = new ParallelTestRunner(globalConfig, context);
        this.sequentialRunner = new SequentialTestRunner({ ...globalConfig, maxWorkers: 1 }, context);
    }

    on<Name extends keyof TestEvents>(
        eventName: Name,
        listener: (eventData: TestEvents[Name]) => void | Promise<void>,
    ): UnsubscribeFn {
        const parallelUnsubscribe = this.parallelRunner.on(eventName, listener);
        const sequentialUnsubscribe = this.sequentialRunner.on(eventName, listener);

        return () => {
            parallelUnsubscribe();
            sequentialUnsubscribe();
        };
    }

    async runTests(tests: Test[], watcher: TestWatcher): Promise<void> {
        await Promise.all([
            this.sequentialRunner.runTests(tests, watcher),
            this.parallelRunner.runTests(tests, watcher),
        ]);
    }
}

class ParallelTestRunner extends TestRunner {
    async runTests(tests: Test[], watcher: TestWatcher): Promise<void> {
        const isParallelizeable = (test: Test) => !test.path.includes(".sequential.");
        const parallelizeableTests = tests.filter(isParallelizeable);

        await super.runTests(parallelizeableTests, watcher, { serial: false });
    }
}

class SequentialTestRunner extends TestRunner {
    async runTests(tests: Test[], watcher: TestWatcher): Promise<void> {
        const isSequentialTest = (test: Test) => test.path.includes(".sequential.");
        const sequential = tests.filter(isSequentialTest);

        // { serial: false } triggers the parallel test execution leveraging workers,
        // but along with { maxWorkers: 1, workerIdleMemoryLimit: ... } configuration
        // this sole worker would be killed and restarted if it exceeds the memory threshold.
        await super.runTests(sequential, watcher, { serial: false });
    }
}