// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import assert from 'assert';
import * as fs from 'fs';
import { InferenceSession, Tensor, TypedTensor } from 'onnxruntime-common';
import * as path from 'path';

import { assertTensorEqual } from '../../test-utils';

const SQUEEZENET_INPUT0_DATA = require(path.join(__dirname, '../../testdata/squeezenet.input0.json'));
const SQUEEZENET_OUTPUT0_DATA = require(path.join(__dirname, '../../testdata/squeezenet.output0.json'));

describe('UnitTests - InferenceSession.create()', () => {
  const modelPath = path.join(__dirname, '../../testdata/squeezenet.onnx');
  const modelBuffer = fs.readFileSync(modelPath);
  const createAny: any = InferenceSession.create;

  // #region test bad arguments
  it('BAD CALL - no argument', async () => {
    await assert.rejects(
      async () => {
        await createAny();
      },
      { name: 'TypeError', message: /argument\[0\]/ },
    );
  });
  it('BAD CALL - byteOffset negative number (ArrayBuffer, number)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer.buffer, -1);
      },
      { name: 'RangeError', message: /'byteOffset'/ },
    );
  });
  it('BAD CALL - byteOffset out of range (ArrayBuffer, number)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer.buffer, 100000000);
      },
      { name: 'RangeError', message: /'byteOffset'/ },
    );
  });
  it('BAD CALL - byteLength negative number (ArrayBuffer, number)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer.buffer, 0, -1);
      },
      { name: 'RangeError', message: /'byteLength'/ },
    );
  });
  it('BAD CALL - byteLength out of range (ArrayBuffer, number)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer.buffer, 0, 100000000);
      },
      { name: 'RangeError', message: /'byteLength'/ },
    );
  });
  it('BAD CALL - options type mismatch (string, string)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelPath, 'cpu');
      },
      { name: 'TypeError', message: /'options'/ },
    );
  });
  it('BAD CALL - options type mismatch (Uint8Array, string)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer, 'cpu');
      },
      { name: 'TypeError', message: /'options'/ },
    );
  });
  it('BAD CALL - options type mismatch (ArrayBuffer, number, number, string)', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelBuffer.buffer, modelBuffer.byteOffset, modelBuffer.byteLength, 'cpu');
      },
      { name: 'TypeError', message: /'options'/ },
    );
  });

  it('EXPECTED FAILURE - Load model failed', async () => {
    await assert.rejects(
      async () => {
        await InferenceSession.create('/this/is/an/invalid/path.onnx');
      },
      { name: 'Error', message: /failed/ },
    );
  });
  it('EXPECTED FAILURE - empty buffer', async () => {
    await assert.rejects(
      async () => {
        await InferenceSession.create(new Uint8Array(0));
      },
      { name: 'Error', message: /No graph was found in the protobuf/ },
    );
  });
  // #endregion

  it('metadata: inputNames', async () => {
    const session = await InferenceSession.create(modelPath);
    assert.deepStrictEqual(session.inputNames, ['data_0']);
  });
  it('metadata: outputNames', async () => {
    const session = await InferenceSession.create(modelPath);
    assert.deepStrictEqual(session.outputNames, ['softmaxout_1']);
  });
});

describe('UnitTests - InferenceSession.run()', () => {
  let session: InferenceSession | null = null;
  let sessionAny: any;
  const input0 = new Tensor('float32', SQUEEZENET_INPUT0_DATA, [1, 3, 224, 224]);
  const expectedOutput0 = new Tensor('float32', SQUEEZENET_OUTPUT0_DATA, [1, 1000, 1, 1]);

  before(async () => {
    session = await InferenceSession.create(path.join(__dirname, '../../testdata/squeezenet.onnx'));
    sessionAny = session;
  });

  // #region test bad input(feeds)
  it('BAD CALL - input type mismatch (null)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run(null);
      },
      { name: 'TypeError', message: /'feeds'/ },
    );
  });
  it('BAD CALL - input type mismatch (single tensor)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run(input0);
      },
      { name: 'TypeError', message: /'feeds'/ },
    );
  });
  it('BAD CALL - input type mismatch (tensor array)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run([input0]);
      },
      { name: 'TypeError', message: /'feeds'/ },
    );
  });
  it('EXPECTED FAILURE - input name missing', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({});
      },
      { name: 'Error', message: /input 'data_0' is missing/ },
    );
  });
  it('EXPECTED FAILURE - input name incorrect', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_1: input0 }); // correct name should be 'data_0'
      },
      { name: 'Error', message: /input 'data_0' is missing/ },
    );
  });
  // #endregion

  // #region test fetches overrides
  it('run() - no fetches', async () => {
    const result = await session!.run({ data_0: input0 });
    assertTensorEqual(result.softmaxout_1, expectedOutput0);
  });
  it('run() - fetches names', async () => {
    const result = await session!.run({ data_0: input0 }, ['softmaxout_1']);
    assertTensorEqual(result.softmaxout_1, expectedOutput0);
  });
  it('run() - fetches object', async () => {
    const result = await session!.run({ data_0: input0 }, { softmaxout_1: null });
    assertTensorEqual(result.softmaxout_1, expectedOutput0);
  });
  // TODO: enable after buffer reuse is implemented
  it.skip('run() - fetches object (pre-allocated)', async () => {
    const preAllocatedOutputBuffer = new Float32Array(expectedOutput0.size);
    const result = await session!.run(
      { data_0: input0 },
      { softmaxout_1: new Tensor(preAllocatedOutputBuffer, expectedOutput0.dims) },
    );
    const softmaxout_1 = result.softmaxout_1 as TypedTensor<'float32'>;
    assert.strictEqual(softmaxout_1.data.buffer, preAllocatedOutputBuffer.buffer);
    assert.strictEqual(softmaxout_1.data.byteOffset, preAllocatedOutputBuffer.byteOffset);
    assertTensorEqual(result.softmaxout_1, expectedOutput0);
  });
  // #endregion

  // #region test bad output(fetches)
  it('BAD CALL - fetches type mismatch (null)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, null);
      },
      { name: 'TypeError', message: /argument\[1\]/ },
    );
  });
  it('BAD CALL - fetches type mismatch (number)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, 1);
      },
      { name: 'TypeError', message: /argument\[1\]/ },
    );
  });
  it('BAD CALL - fetches type mismatch (Tensor)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run(
          { data_0: input0 },
          new Tensor(new Float32Array(expectedOutput0.size), expectedOutput0.dims),
        );
      },
      { name: 'TypeError', message: /'fetches'/ },
    );
  });
  it('BAD CALL - fetches as array (empty array)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, []);
      },
      { name: 'TypeError', message: /'fetches'/ },
    );
  });
  it('BAD CALL - fetches as array (non-string elements)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, [1, 2, 3]);
      },
      { name: 'TypeError', message: /'fetches'/ },
    );
  });
  it('BAD CALL - fetches as array (invalid name)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, ['im_a_wrong_output_name']);
      },
      { name: 'RangeError', message: /'fetches'/ },
    );
  });
  // #endregion

  it('BAD CALL - options type mismatch (number)', async () => {
    await assert.rejects(
      async () => {
        await sessionAny.run({ data_0: input0 }, ['softmaxout_1'], 1);
      },
      { name: 'TypeError', message: /'options'/ },
    );
  });
});

describe('UnitTests - InferenceSession.SessionOptions', () => {
  const modelPath = path.join(__dirname, '../../testdata/test_types_float.onnx');
  const createAny: any = InferenceSession.create;

  it('BAD CALL - type mismatch', async () => {
    await assert.rejects(
      async () => {
        await createAny(modelPath, 'cpu');
      },
      { name: 'TypeError', message: /'options'/ },
    );
  });

  describe('executionProviders', () => {
    it.skip('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { executionProviders: 'bad-EP-name' });
        },
        { name: 'TypeError', message: /executionProviders/ },
      );
    });
    it.skip('EXPECTED FAILURE - invalid EP name, string list', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { executionProviders: ['bad-EP-name'] });
        },
        { name: 'Error', message: /executionProviders.+bad-EP-name/ },
      );
    });
    it.skip('EXPECTED FAILURE - invalid EP name, object list', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { executionProviders: [{ name: 'bad-EP-name' }] });
        },
        { name: 'Error', message: /executionProviders.+bad-EP-name/ },
      );
    });
    it('string list (CPU)', async () => {
      await InferenceSession.create(modelPath, { executionProviders: ['cpu'] });
    });
    it('object list (CPU)', async () => {
      await InferenceSession.create(modelPath, { executionProviders: [{ name: 'cpu' }] });
    });
  });

  describe('intraOpNumThreads', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { intraOpNumThreads: 'bad-value' });
        },
        { name: 'TypeError', message: /intraOpNumThreads/ },
      );
    });
    it('BAD CALL - non-integer', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { intraOpNumThreads: 1.5 });
        },
        { name: 'RangeError', message: /intraOpNumThreads/ },
      );
    });
    it('BAD CALL - negative integer', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { intraOpNumThreads: -1 });
        },
        { name: 'RangeError', message: /intraOpNumThreads/ },
      );
    });
    it('intraOpNumThreads = 1', async () => {
      await InferenceSession.create(modelPath, { intraOpNumThreads: 1 });
    });
  });

  describe('interOpNumThreads', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { interOpNumThreads: 'bad-value' });
        },
        { name: 'TypeError', message: /interOpNumThreads/ },
      );
    });
    it('BAD CALL - non-integer', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { interOpNumThreads: 1.5 });
        },
        { name: 'RangeError', message: /interOpNumThreads/ },
      );
    });
    it('BAD CALL - negative integer', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { interOpNumThreads: -1 });
        },
        { name: 'RangeError', message: /interOpNumThreads/ },
      );
    });
    it('interOpNumThreads = 1', async () => {
      await InferenceSession.create(modelPath, { interOpNumThreads: 1 });
    });
  });

  describe('graphOptimizationLevel', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { graphOptimizationLevel: 0 });
        },
        { name: 'TypeError', message: /graphOptimizationLevel/ },
      );
    });
    it('BAD CALL - invalid config', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { graphOptimizationLevel: 'bad-value' });
        },
        { name: 'TypeError', message: /graphOptimizationLevel/ },
      );
    });
    it('graphOptimizationLevel = basic', async () => {
      await InferenceSession.create(modelPath, { graphOptimizationLevel: 'basic' });
    });
  });

  describe('enableCpuMemArena', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { enableCpuMemArena: 0 });
        },
        { name: 'TypeError', message: /enableCpuMemArena/ },
      );
    });
    it('enableCpuMemArena = true', async () => {
      await InferenceSession.create(modelPath, { enableCpuMemArena: true });
    });
  });

  describe('enableMemPattern', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { enableMemPattern: 0 });
        },
        { name: 'TypeError', message: /enableMemPattern/ },
      );
    });
    it('enableMemPattern = true', async () => {
      await InferenceSession.create(modelPath, { enableMemPattern: true });
    });
  });

  describe('executionMode', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { executionMode: 0 });
        },
        { name: 'TypeError', message: /executionMode/ },
      );
    });
    it('BAD CALL - invalid config', async () => {
      await assert.rejects(
        async () => {
          await createAny(modelPath, { executionMode: 'bad-value' });
        },
        { name: 'TypeError', message: /executionMode/ },
      );
    });
    it('executionMode = sequential', async () => {
      await InferenceSession.create(modelPath, { executionMode: 'sequential' });
    });
  });
});

describe('UnitTests - InferenceSession.RunOptions', () => {
  let session: InferenceSession | null = null;
  let sessionAny: any;
  const input0 = new Tensor('float32', [1, 2, 3, 4, 5], [1, 5]);
  const expectedOutput0 = new Tensor('float32', [1, 2, 3, 4, 5], [1, 5]);

  before(async () => {
    const modelPath = path.join(__dirname, '../../testdata/test_types_float.onnx');
    session = await InferenceSession.create(modelPath);
    sessionAny = session;
  });

  describe('logSeverityLevel', () => {
    it('BAD CALL - type mismatch', async () => {
      await assert.rejects(
        async () => {
          await sessionAny.run({ input: input0 }, { logSeverityLevel: 'error' });
        },
        { name: 'TypeError', message: /logSeverityLevel/ },
      );
    });
    it('BAD CALL - out of range', async () => {
      await assert.rejects(
        async () => {
          await sessionAny.run({ input: input0 }, { logSeverityLevel: 8 });
        },
        { name: 'RangeError', message: /logSeverityLevel/ },
      );
    });
    it('BAD CALL - out of range', async () => {
      await assert.rejects(
        async () => {
          await sessionAny.run({ input: input0 }, { logSeverityLevel: 8 });
        },
        { name: 'RangeError', message: /logSeverityLevel/ },
      );
    });
    it('logSeverityLevel = 4', async () => {
      const result = await sessionAny.run({ input: input0 }, { logSeverityLevel: 4 });
      assertTensorEqual(result.output, expectedOutput0);
    });
  });
});