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

import * as assert from 'assert';
import * as fs from 'fs';
import { EOL } from 'os';
import * as path from 'path';

import { Attribute } from '../lib/onnxjs/attribute';
import { WEBGL_OP_RESOLVE_RULES } from '../lib/onnxjs/backends/webgl/op-resolve-rules';
import { OpSet, resolveOperator } from '../lib/onnxjs/opset';
import { Tensor } from '../lib/onnxjs/tensor';

function checkSupport(type: string, range: [number, number], rules: readonly OpSet.ResolveRule[]) {
  const node = { name: '', opType: type, inputs: [], outputs: [], attributes: new Attribute(undefined) };
  for (let i = range[0]; i <= range[1]; i++) {
    try {
      resolveOperator(node, [{ domain: '', version: i }], rules);
    } catch (_e) {
      return false;
    }
  }
  return true;
}

function formatDesc(opType: string, range: [number, number], support: boolean, last: boolean) {
  let versionDesc = '';
  if (support) {
    versionDesc = last ? `${range[0]}+` : range[0] === range[1] ? `${range[0]}` : `${range[0]}-${range[1]}`;
    versionDesc = `[${versionDesc}](https://github.com/onnx/onnx/blob/main/docs/Changelog.md#${opType}-${range[0]})`;
  }
  return versionDesc;
}
function dummyOpImpl(): Tensor[] {
  return {} as any as Tensor[];
}

const ops = new Map<string, Map<string, number[]>>();
const webglCheckOnlyRules = WEBGL_OP_RESOLVE_RULES.map(
  (rule) => [rule[0], rule[1], rule[2], dummyOpImpl] as OpSet.ResolveRule,
);

fs.readFileSync(path.join(__dirname, '../../../cmake/external/onnx/onnx/defs/operator_sets.h'), 'utf8')
  .split(/\r?\n/)
  .forEach((line) => {
    const matcher = /class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME\(\s*(\w+),\s*(\d+),\s*(\w+)\)/;
    const matches = matcher.exec(line);
    if (matches) {
      const opset = matches[1];
      const version = Number.parseInt(matches[2], 10);
      const opType = matches[3];

      let currentSet = ops.get(opset);
      if (currentSet === undefined) {
        currentSet = new Map<string, number[]>();
        ops.set(opset, currentSet);
      }

      let currentOp = currentSet.get(opType);
      if (currentOp === undefined) {
        currentOp = [];
        currentSet.set(opType, currentOp);
      }

      currentOp.push(version);
    }
  });

const opsets = Array.from(ops.keys());
assert.ok(opsets.length === 1 && opsets[0] === 'Onnx');

const onnxOpset = ops.get(opsets[0])!;
const opTypes = Array.from(onnxOpset.keys()).sort();

const doc = fs.createWriteStream(path.join(__dirname, '../docs/webgl-operators.md'));
doc.write(`## Operators Support Table${EOL}${EOL}`);
doc.write(`The following table shows [ai.onnx](https://github.com/onnx/onnx/blob/main/docs/Operators.md)\
  operators from which onnx opset version are currently supported by ONNX Runtime Web. For example, \`4-6, 8+\` means\
  ONNX Runtime Web currently support opset version 4 to 6, 8 and above.${EOL}${EOL}`);
doc.write(`See [Compatibility](../README.md#Compatibility) for a list of the supported platforms.${EOL}${EOL}`);
doc.write(`*This file is automatically generated from the\
  def files via [this script](../script/generate-operator-md.ts).\
  Do not modify directly.*${EOL}${EOL}`);
doc.write(`| Operator | WebGl Backend |${EOL}`);
doc.write(`|:--------:|:-------------:|${EOL}`);

let VERSION_MAX = 0;
onnxOpset.forEach((versions) => {
  versions.forEach((version) => (VERSION_MAX = Math.max(VERSION_MAX, version)));
});

for (const type of opTypes) {
  const versions = onnxOpset.get(type)!.sort((a, b) => a - b);

  const webgl: string[] = [];
  for (let i = 0; i < versions.length; i++) {
    const last = i === versions.length - 1;
    const versionRange: [number, number] = [versions[i], last ? VERSION_MAX : versions[i + 1] - 1];

    webgl.push(formatDesc(type, versionRange, checkSupport(type, versionRange, webglCheckOnlyRules), last));
  }

  doc.write(
    `| [${type}](https://github.com/onnx/onnx/blob/main/docs/Operators.md#${type}) | ${webgl
      .filter((d) => d.length > 0)
      .join(', ')} |${EOL}`,
  );
}
doc.end();
