import {
  CliLeaf,
  CliNumberInput,
  CliOneOfInput,
  CliStringInput,
  CliTerseError
} from '@alwaysai/alwayscli';
import * as logSymbols from 'log-symbols';
import { join } from 'path';
import * as tempy from 'tempy';
import { checkUserIsLoggedInComponent } from '../../components/user';
import { appInstallModel, getTargetHardwareType } from '../../core/app';
import {
  downloadModelPackageToCache,
  ModelId,
  modelPackageCache
} from '../../core/model';
import { requirePaidPlan } from '../../core/project';
import { CliRpcClient } from '../../infrastructure';
import {
  DockerSpawner,
  echo,
  JsSpawner,
  keyMirror,
  logger,
  stringifyError
} from '../../util';

export const OutputFormat = keyMirror({
  'tensor-rt': null
  // hailo: null,
});

export type OutputFormat = keyof typeof OutputFormat;
export const OUTPUT_FORMAT_OPTIONS = Object.values(OutputFormat);

export const modelConvert = CliLeaf({
  name: 'convert',
  description: 'Convert an alwaysAI model to a different format',
  positionalInput: CliStringInput({
    placeholder: 'model ID e.g. alwaysai/mobilenet_ssd',
    required: true
  }),
  namedInputs: {
    format: CliOneOfInput({
      description: 'The output format of the model conversion',
      required: true,
      values: OUTPUT_FORMAT_OPTIONS
    }),
    'output-id': CliStringInput({
      description: 'Model ID of the converted model',
      required: false
    }),
    version: CliNumberInput({
      description: 'Version of the model to convert',
      required: false
    }),
    'batch-size': CliNumberInput({
      description: 'Batch size if converting to tensor-rt',
      required: false
    })
  },
  async action(
    modelId,
    { format, 'output-id': outputId, version, 'batch-size': batchSize }
  ) {
    await convertModel({ modelId, format, outputId, version, batchSize });
  }
});

async function convertModel(opts: {
  modelId: string;
  format: OutputFormat;
  outputId?: string;
  version?: number;
  batchSize?: number;
}) {
  const { modelId, format, outputId, batchSize } = opts;
  await checkUserIsLoggedInComponent({ yes: true });

  await requirePaidPlan();

  let version = opts.version;
  try {
    const modelDetails = await CliRpcClient().getModelVersion({ id: modelId });
    if (version === undefined) {
      version = modelDetails.version;
    }
    if (!modelPackageCache.has(modelId, version)) {
      await downloadModelPackageToCache(modelId, version);
    }
  } catch (err) {
    logger.error(
      `Model ${modelId}: ${version} not found: ${stringifyError(err)}`
    );
    throw new CliTerseError(`Model ${modelId}: ${version} not found!`);
  }

  if (batchSize && format !== 'tensor-rt') {
    throw new CliTerseError(
      `Batch size parameter is only supported for tensor-rt conversion. Please remove the batch size parameter.`
    );
  }

  if (outputId) {
    ModelId.parse(outputId);
  }

  switch (format) {
    case 'tensor-rt':
      await convertTensorrt({ modelId, version, outputId, batchSize });
      break;

    // case 'hailo':
    // await convertHailo(modelId, outputId); //disable hailo conversion
    // break;

    default:
      throw new CliTerseError(
        `${format} is not a valid format. Choose one from the following: <${OUTPUT_FORMAT_OPTIONS}>`
      );
  }
}

// async function convertHailo(modelId: string, outputId: string) {
//   echo('Starting conversion to Hailo...');
//   const hailoDockerImage = 'tonyjesudoss/hailomodelconvertor'; // will update with new image once released.

//   await JsSpawner().runForeground({
//     exe: 'docker',
//     args: [
//       'run',
//       '--rm',
//       '-it',
//       '--env',
//       `MODEL_ID=${modelId}`,
//       '--env',
//       `OUTPUT_MODEL_ID=${outputId}`,
//       hailoDockerImage,
//     ],
//   });
// }

async function convertTensorrt(props: {
  modelId: string;
  version: number;
  outputId?: string;
  batchSize?: number;
}) {
  echo('Starting conversion to TensorRT...');
  const { modelId, version, outputId, batchSize } = props;

  // TODO: enable conversion on remote device
  const targetHardware = await getTargetHardwareType({});
  if (!targetHardware.includes('jetson')) {
    throw new CliTerseError('NVIDIA Jetson required to convert to TensorRT');
  }

  const tmpDir = tempy.directory();

  const tmpSpawner = JsSpawner({ path: tmpDir });
  await tmpSpawner.mkdirp();

  let bashArgs = `python3 app.py --model-id ${modelId}`;
  if (outputId) {
    bashArgs = `${bashArgs} --output-id ${outputId}`;
  }
  if (batchSize) {
    bashArgs = `${bashArgs} --batch-size ${batchSize}`;
  }

  try {
    await appInstallModel(tmpSpawner, modelId, version);

    const dockerImageId = `alwaysai/model-conversion:tensorrt-${targetHardware}`;
    await JsSpawner().runForeground({
      exe: 'docker',
      args: ['pull', dockerImageId]
    });

    const result = await DockerSpawner({
      dockerImageId,
      targetHardware,
      volumes: [
        `${join(tmpDir, 'models')}:/convert/models`,
        `${join(process.cwd(), 'out')}:/convert/out`
      ]
    }).runForeground({
      cwd: '/convert',
      exe: 'bash',
      args: ['-c', `${bashArgs}`],
      superuser: true
    });
    if (result !== undefined) {
      logger.error(`Model conversion failed with error code ${result}`);
      throw new Error(`Model conversion failed with error code ${result}`);
    }

    if (outputId) {
      const newId = ModelId.parse(outputId);
      echo(
        `${logSymbols.success} Converted model saved to ./out/${newId.publisher}/${newId.name}`
      );
    } else {
      echo(`${logSymbols.success} Converted model saved to ./out/`);
    }
  } catch (err) {
    logger.error(stringifyError(err));
    throw new CliTerseError(
      'Model conversion failed! See errors in above logs.'
    );
  } finally {
    // Clean up temp directory
    await tmpSpawner.rimraf();
  }
}