import { stringify } from 'yaml';
import { TargetHardware } from '../../core/app';

export type DockerComposeRestart =
  | 'no'
  | 'always'
  | 'on-failure'
  | 'unless-stopped';

export type DockerComposeCmd = {
  build: {
    context: string;
    dockerfile?: string;
    args?: object;
  };
  privileged?: boolean;
  volumes?: string[];
  user?: string;
  ports?: string[];
  targetHardware?: TargetHardware;
  workdir?: string;
  restart?: DockerComposeRestart;
};

export function getDockerComposeFileContents(cmd: DockerComposeCmd): string {
  const alwaysai = {
    build: cmd.build
  };
  if (cmd.privileged) {
    alwaysai['privileged'] = cmd.privileged;
  }
  if (cmd.user) {
    alwaysai['user'] = cmd.user;
  }
  if (cmd.ports) {
    alwaysai['ports'] = cmd.ports;
  } else {
    alwaysai['network_mode'] = 'host';
  }
  if (
    cmd.targetHardware?.includes('jetson') ||
    cmd.targetHardware?.includes('x86-trt')
  ) {
    alwaysai['runtime'] = 'nvidia';
  }
  if (cmd.targetHardware?.includes('jetson')) {
    alwaysai['ipc'] = 'host';
    if (cmd.volumes) {
      cmd.volumes.push('/tmp/argus_socket:/tmp/argus_socket');
    } else {
      cmd.volumes = ['/tmp/argus_socket:/tmp/argus_socket'];
    }
  }
  if (cmd.volumes) {
    alwaysai['volumes'] = cmd.volumes;
  }
  if (cmd.workdir) {
    alwaysai['working_dir'] = cmd.workdir;
  }
  if (cmd.restart) {
    alwaysai['restart'] = cmd.restart;
  }
  return stringify({ services: { alwaysai } });
}