import { join } from 'path';
import { parse, stringify } from 'yaml';
import { DOCKER_COMPOSE_FILE } from '../../paths';
import { JsSpawner } from '../../util';
import { TargetHardware } from './get-target-hardware-type';

/**
 * Updates `docker-compose.yaml` based on the target hardware
 *
 * Since it is not known in advance which services contain alwaysAI apps
 * (although the default alwaysAI service is called "alwaysai") any service with
 * a build section will be updated based on target hardware changes. Changes
 * include the BUILD_ARG, runtime, and volumes.
 *
 * If the target hardware configured in the `docker-compose.yaml` file is the
 * same as the new one, no updates are made.
 */
export async function updateDockerComposeTargetHw(opts: {
  targetHardware: TargetHardware;
  cwd?: string;
}) {
  const { targetHardware } = opts;
  const path = opts.cwd
    ? join(opts.cwd, DOCKER_COMPOSE_FILE)
    : DOCKER_COMPOSE_FILE;
  const spawner = JsSpawner();
  if (!(await spawner.exists(path))) {
    throw Error(`${path} not found!`);
  }
  const serialized = await spawner.readFile(path);
  const contents = parse(serialized);
  let changed = false;
  for (const serviceName in contents['services']) {
    const service = contents['services'][serviceName];
    if (!('build' in service)) {
      // Only consider services with a "build" section, since all alwaysAI
      // services are expected to have it.
      continue;
    }
    if (!('args' in service['build'])) {
      contents['services'][serviceName]['build']['args'] = {};
    }
    const args = service['build']['args'];
    if ('ALWAYSAI_HW' in args && args['ALWAYSAI_HW'] === targetHardware) {
      // Skip further updates since target hardware is not changing
      continue;
    }
    contents['services'][serviceName]['build']['args']['ALWAYSAI_HW'] =
      targetHardware;

    // For NVIDIA devices we need to update a few more configurations.
    // NOTE: This doesn't handle the case where the file was based on an NVIDIA
    // device but is now on a non-NVIDIA device.
    if (
      targetHardware.includes('jetson') ||
      targetHardware.includes('x86-trt')
    ) {
      contents['services'][serviceName]['runtime'] = 'nvidia';
    }
    if (targetHardware.includes('jetson')) {
      contents['services'][serviceName]['ipc'] = 'host';
      const argusSocketMount = '/tmp/argus_socket:/tmp/argus_socket';
      if ('volumes' in service) {
        let argusSocketExists = false;
        for (const volume of service['volumes']) {
          if (volume === argusSocketMount) {
            argusSocketExists = true;
          }
        }
        if (!argusSocketExists) {
          contents['services'][serviceName]['volumes'].push(argusSocketMount);
        }
      } else {
        contents['services'][serviceName]['volumes'] = [argusSocketMount];
      }
    }
    changed = true;
  }
  if (changed === true) {
    const outSerialized = stringify(contents);
    await spawner.writeFile(path, outSerialized);
  }
}