// Copyright (c) 2017, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

import 'dart:async';
import 'dart:collection';
import 'dart:io';

import 'package:bazel_worker/bazel_worker.dart';
import 'package:bazel_worker/driver.dart';
import 'package:test/test.dart';

void main() {
  BazelWorkerDriver? driver;
  final disconnectedResponse = WorkResponse(
    exitCode: EXIT_CODE_BROKEN_PIPE,
    output: 'Connection closed',
  )..freeze();

  group('basic driver', () {
    test('can run a single request', () async {
      await _doRequests(count: 1);
      await _doRequests(count: 1);
    });

    test('can run multiple batches of requests through multiple workers',
        () async {
      var maxWorkers = 4;
      var maxIdleWorkers = 2;
      driver = BazelWorkerDriver(MockWorker.spawn,
          maxWorkers: maxWorkers, maxIdleWorkers: maxIdleWorkers);
      for (var i = 0; i < 10; i++) {
        await _doRequests(driver: driver);
        expect(MockWorker.liveWorkers.length, maxIdleWorkers);
        // No workers should be killed while there is ongoing work, but they
        // should be cleaned up once there isn't any more work to do.
        expect(MockWorker.deadWorkers.length,
            (maxWorkers - maxIdleWorkers) * (i + 1));
      }
    });

    test('can run multiple requests through one worker', () async {
      var maxWorkers = 1;
      var maxIdleWorkers = 1;
      driver = BazelWorkerDriver(MockWorker.spawn,
          maxWorkers: maxWorkers, maxIdleWorkers: maxIdleWorkers);
      for (var i = 0; i < 10; i++) {
        await _doRequests(driver: driver);
        expect(MockWorker.liveWorkers.length, 1);
        expect(MockWorker.deadWorkers.length, 0);
      }
    });

    test('can run one request through multiple workers', () async {
      driver =
          BazelWorkerDriver(MockWorker.spawn, maxWorkers: 4, maxIdleWorkers: 4);
      for (var i = 0; i < 10; i++) {
        await _doRequests(driver: driver, count: 1);
        expect(MockWorker.liveWorkers.length, 1);
        expect(MockWorker.deadWorkers.length, 0);
      }
    });

    test('can run with maxIdleWorkers == 0', () async {
      var maxWorkers = 4;
      driver = BazelWorkerDriver(MockWorker.spawn,
          maxWorkers: maxWorkers, maxIdleWorkers: 0);
      for (var i = 0; i < 10; i++) {
        await _doRequests(driver: driver);
        expect(MockWorker.liveWorkers.length, 0);
        expect(MockWorker.deadWorkers.length, maxWorkers * (i + 1));
      }
    });

    test('trackWork gets invoked when a worker is actually ready', () async {
      var maxWorkers = 2;
      driver = BazelWorkerDriver(MockWorker.spawn, maxWorkers: maxWorkers);
      var tracking = <Future>[];
      await _doRequests(
          driver: driver,
          count: 10,
          trackWork: (Future response) {
            // We should never be tracking more than `maxWorkers` jobs at a time.
            expect(tracking.length, lessThan(maxWorkers));
            tracking.add(response);
            response.then((_) => tracking.remove(response));
          });
    });

    group('failing workers', () {
      /// A driver which spawns [numBadWorkers] failing workers and then good
      /// ones after that, and which will retry [maxRetries] times.
      void createDriver({int maxRetries = 2, int numBadWorkers = 2}) {
        var numSpawned = 0;
        driver = BazelWorkerDriver(
            () async => MockWorker(workerLoopFactory: (MockWorker worker) {
                  var connection = StdAsyncWorkerConnection(
                      inputStream: worker._stdinController.stream,
                      outputStream: worker._stdoutController.sink);
                  if (numSpawned < numBadWorkers) {
                    numSpawned++;
                    return ThrowingMockWorkerLoop(
                        worker, MockWorker.responseQueue, connection);
                  } else {
                    return MockWorkerLoop(MockWorker.responseQueue,
                        connection: connection);
                  }
                }),
            maxRetries: maxRetries);
      }

      test('should retry up to maxRetries times', () async {
        createDriver();
        var expectedResponse = WorkResponse();
        MockWorker.responseQueue.addAll(
            [disconnectedResponse, disconnectedResponse, expectedResponse]);
        var actualResponse = await driver!.doWork(WorkRequest());
        // The first 2 null responses are thrown away, and we should get the
        // third one.
        expect(actualResponse, expectedResponse);

        expect(MockWorker.deadWorkers.length, 2);
        expect(MockWorker.liveWorkers.length, 1);
      });

      test('should fail if it exceeds maxRetries failures', () async {
        createDriver(maxRetries: 2, numBadWorkers: 3);
        MockWorker.responseQueue.addAll(
            [disconnectedResponse, disconnectedResponse, WorkResponse()]);
        var actualResponse = await driver!.doWork(WorkRequest());
        // Should actually get a bad response.
        expect(actualResponse.exitCode, 15);
        expect(
            actualResponse.output,
            'Invalid response from worker, this probably means it wrote '
            'invalid output or died.');

        expect(MockWorker.deadWorkers.length, 3);
      });
    });

    test('handles spawnWorker failures', () async {
      driver = BazelWorkerDriver(() async => throw StateError('oh no!'),
          maxRetries: 0);
      expect(driver!.doWork(WorkRequest()), throwsA(isA<StateError>()));
    });

    tearDown(() async {
      await driver?.terminateWorkers();
      expect(MockWorker.liveWorkers, isEmpty);
      MockWorker.deadWorkers.clear();
      MockWorker.responseQueue.clear();
    });
  });
}

/// Runs [count] of fake work requests through [driver], and asserts that they
/// all completed.
Future _doRequests(
    {BazelWorkerDriver? driver,
    int count = 100,
    void Function(Future<WorkResponse?>)? trackWork}) async {
  // If we create a driver, we need to make sure and terminate it.
  var terminateDriver = driver == null;
  driver ??= BazelWorkerDriver(MockWorker.spawn);
  var requests = List.generate(count, (_) => WorkRequest());
  var responses = List.generate(count, (_) => WorkResponse());
  MockWorker.responseQueue.addAll(responses);
  var actualResponses = await Future.wait(
      requests.map((request) => driver!.doWork(request, trackWork: trackWork)));
  expect(actualResponses, unorderedEquals(responses));
  if (terminateDriver) await driver.terminateWorkers();
}

/// A mock worker loop that returns work responses from the provided list.
///
/// Throws if it runs out of responses.
class MockWorkerLoop extends AsyncWorkerLoop {
  final Queue<WorkResponse> _responseQueue;

  MockWorkerLoop(this._responseQueue, {super.connection});

  @override
  Future<WorkResponse> performRequest(WorkRequest request) async {
    print('Performing request $request');
    return _responseQueue.removeFirst();
  }
}

/// A mock worker loop with a custom `run` function that throws.
class ThrowingMockWorkerLoop extends MockWorkerLoop {
  final MockWorker _mockWorker;

  ThrowingMockWorkerLoop(this._mockWorker, Queue<WorkResponse> responseQueue,
      AsyncWorkerConnection connection)
      : super(responseQueue, connection: connection);

  /// Run the worker loop. The returned [Future] doesn't complete until
  /// [connection#readRequest] returns `null`.
  @override
  Future run() async {
    while (true) {
      var request = await connection.readRequest();
      if (request == null) break;
      await performRequest(request);
      _mockWorker.kill();
    }
  }
}

/// A mock worker process.
///
/// Items in [responseQueue] will be returned in order based on requests.
///
/// If there are no items left in [responseQueue] then it will throw.
class MockWorker implements Process {
  /// Spawns a new [MockWorker].
  static Future<MockWorker> spawn() async => MockWorker();

  /// Static queue of pending responses, these are shared by all workers.
  ///
  /// If this is empty and a request is received then it will throw.
  static final responseQueue = Queue<WorkResponse>();

  /// Static list of all live workers.
  static final liveWorkers = <MockWorker>[];

  /// Static list of all the dead workers.
  static final deadWorkers = <MockWorker>[];

  /// Standard constructor, creates a [WorkerLoop] from [workerLoopFactory] or
  /// a [MockWorkerLoop] if no factory is provided.
  MockWorker({WorkerLoop Function(MockWorker)? workerLoopFactory}) {
    liveWorkers.add(this);
    var workerLoop = workerLoopFactory != null
        ? workerLoopFactory(this)
        : MockWorkerLoop(responseQueue,
            connection: StdAsyncWorkerConnection(
                inputStream: _stdinController.stream,
                outputStream: _stdoutController.sink));
    workerLoop.run();
  }

  @override
  Future<int> get exitCode => throw UnsupportedError('Not needed.');

  @override
  Stream<List<int>> get stdout => _stdoutController.stream;
  final _stdoutController = StreamController<List<int>>();

  @override
  Stream<List<int>> get stderr => _stderrController.stream;
  final _stderrController = StreamController<List<int>>();

  @override
  late final IOSink stdin = IOSink(_stdinController.sink);
  final _stdinController = StreamController<List<int>>();

  @override
  int get pid => throw UnsupportedError('Not needed.');

  @override
  bool kill(
      [ProcessSignal processSignal = ProcessSignal.sigterm, int exitCode = 0]) {
    if (_killed) return false;
    () async {
      await _stdoutController.close();
      await _stderrController.close();
      await _stdinController.close();
    }();
    deadWorkers.add(this);
    liveWorkers.remove(this);
    return true;
  }

  final _killed = false;
}