// Experimental

import * as flatbuffers from './flatbuffers.js';

const megengine = {};

megengine.ModelFactory = class {

    async match(context) {
        const stream = context.stream;
        if (stream && stream.length >= 12) {
            let buffer = stream.peek(12);
            const tag = String.fromCharCode.apply(null, buffer);
            const position = tag.startsWith('mgbtest0') ? 12 : 0;
            if (stream.length > (position + 12)) {
                buffer = stream.peek(24).slice(position, position + 12);
                const size = (buffer[0] | buffer[1] << 8 | buffer[2] << 16 | buffer[3] << 24) >>> 0;
                if (position > 0 || size === (stream.length - position - 4)) {
                    const reader = flatbuffers.BinaryReader.open(stream, position + 4);
                    if (reader.identifier === 'mgv2') {
                        return context.set('megengine.mge', reader);
                    }
                }
            }
            for (const value of ['mgb0001', 'mgb0000a', 'MGBS', 'MGBC']) {
                if (tag.startsWith(value)) {
                    return context.set(`megengine.${value}`);
                }
            }
        }
        const obj = await context.peek('pkl');
        if (obj && obj.__class__ && obj.__class__.__module__ === 'megengine.traced_module.traced_module' && obj.__class__.__name__ === 'TracedModule') {
            return context.set('megengine.tm');
        }
        return null;
    }

    async open(context) {
        const metadata = await context.metadata('megengine-metadata.json');
        switch (context.type) {
            case 'megengine.tm': {
                const obj = await context.peek('pkl');
                return new megengine.Model(metadata, obj, context.type);
            }
            case 'megengine.mge': {
                megengine.schema = await context.require('./megengine-schema');
                megengine.schema = megengine.schema.mgb.serialization.fbs;
                let model = null;
                try {
                    const reader = context.value;
                    model = megengine.schema.v2.Model.create(reader);
                } catch (error) {
                    const message = error && error.message ? error.message : error.toString();
                    throw new megengine.Error(`File format is not megengine.Model (${message.replace(/\.$/, '')}).`);
                }
                return new megengine.Model(metadata, model, context.type);
            }
            default: {
                throw new megengine.Error(`Unsupported MegEngine format '${context.type.replace(/^megengine\./, '')}'.`);
            }
        }
    }
};

megengine.Model = class {

    constructor(metadata, obj, type) {
        this.format = 'MegEngine';
        if (type === 'megengine.tm') {
            this.format += obj.dump_info && obj.dump_info.has('version') ? ` v${obj.dump_info.get('version')}` : '';
        } else if (type === 'megengine.mge') {
            this.format += ` Mge${obj.model_version ? ` v${obj.model_version}` : ''}`;
        }
        this.modules = [new megengine.Graph(metadata, obj)];
    }
};

megengine.Graph = class {

    constructor(metadata, obj) {
        this.name = '';
        this.nodes = [];
        this.inputs = [];
        this.outputs = [];
        const values = new Map();
        values.map = (name, type, tensor, quantization) => {
            type = type || null;
            tensor = tensor || null;
            if (tensor) {
                return new megengine.Value(name, type, tensor, quantization);
            }
            if (!values.has(name)) {
                values.set(name, new megengine.Value(name, type, tensor, quantization));
            } else if ((type && !type.equals(values.get(name).type)) || tensor) {
                throw new megengine.Error(`Duplicate value '${name}'.`);
            }
            return values.get(name);
        };
        const loadGraph = (tmodule, igraph, context, namePrefix, metadata, isRoot) => {
            const expressions = igraph._exprs;
            const getTensorType = (dtype, shape) => {
                dtype = dtype ? dtype.__name__ : null;
                return new megengine.TensorType(dtype, new megengine.TensorShape(shape));
            };
            if (isRoot) {
                for (const node of igraph._inputs) {
                    if (node.__class__.__name__ !== 'ModuleNode') {
                        const type = getTensorType(node._dtype, node._shape);
                        const value = values.map(node._name, type, null);
                        const argument = new megengine.Argument(node._name, [value]);
                        this.inputs.push(argument);
                    }
                }
                for (const node of igraph._outputs) {
                    const type = getTensorType(node._dtype, node._shape);
                    const value = values.map(node._name, type, null);
                    const argument = new megengine.Argument(node._name, [value]);
                    this.outputs.push(argument);
                }
            }
            const parseGetAttr = (module, expression) => {
                let names = expression.name.split('.');
                while (expression.inputs[0].expr.__class__.__name__ === 'GetAttr') {
                    expression = expression.inputs[0].expr;
                    names = expression.name.split('.').concat(names);
                }
                let obj = module;
                for (const name of names) {
                    if (obj.__class__.__name__ === '_ModuleState') {
                        obj = obj.state;
                    }
                    obj = obj[name];
                }
                return obj;
            };
            const parseArgs = (args, kwargs, meta) => {
                const state = {};
                let argIdx = 0;
                const processArgs = (inp, startIdx) => {
                    while (typeof inp === 'string' && inp.indexOf('Tensor') !== -1) {
                        inp = inp.replace('Tensor', `input${startIdx === 0 ? '' : startIdx}`);
                        startIdx += 1;
                    }
                    return [inp, startIdx];
                };
                const formatTreeDef = (obj) => {
                    if (obj.__class__.__name__ !== 'TreeDef' && obj.__class__.__name__ !== 'LeafDef') {
                        throw new megengine.Error(`Invalid argument '${obj.__class__.__name__}'.`);
                    }
                    if (obj.__class__.__name__ === 'TreeDef') {
                        const type = typeof obj.type === 'string' ? obj.type.split('.').slice(-1)[0] : obj.type.__name__;
                        const list = obj.children_defs.map((child) => formatTreeDef(child));
                        switch (type) {
                            case 'tuple': {
                                return `(${list.join(',')})`;
                            }
                            case 'slice': {
                                return list.join(':');
                            }
                            case 'list': {
                                return `[${list.join(',')}]`;
                            }
                            case 'dict': {
                                let content = '';
                                for (let i = 0; i < this.children_defs.length; i++) {
                                    content += `${this.aux_data[i]}:${list[i]}`;
                                }
                                return `{${content}}`;
                            }
                            default: {
                                return `${type}(${list.join(',')})`;
                            }
                        }
                    }
                    if (obj.const_val !== null) {
                        return obj.const_val;
                    } else if (obj.type[0].__module__ !== undefined) {
                        return obj.type[0].__name__;
                    }
                    return 'None';
                };
                let inpIdx = 0;
                for (const arg of args.children_defs) {
                    let name = '';
                    if (meta.attributes === undefined || (meta.attributes.length !== args.children_defs.length && meta.varargs === null)) {
                        name = `arg${argIdx}`;
                    } else if (argIdx < meta.attributes.length) {
                        name = meta.attributes[argIdx].name;
                    } else {
                        name = meta.varargs + (argIdx - meta.attributes.length);
                    }
                    const [value, index] = processArgs(formatTreeDef(arg), inpIdx);
                    state[name] = value;
                    inpIdx = index;
                    argIdx += 1;
                }
                for (let i = 0; i < kwargs.children_defs.length; i++) {
                    const [value, index] = processArgs(formatTreeDef(kwargs.children_defs[i]), inpIdx);
                    state[kwargs.aux_data[i]] = value;
                    inpIdx = index;
                }
                return state;
            };
            const getName = (context, name) => {
                let rst = name;
                while (context.get(rst) !== undefined) {
                    if (rst === context.get(rst)) {
                        return rst;
                    }
                    rst = context.get(rst);
                }
                return rst;
            };
            const getFullName = (prefix, name) => {
                return prefix === '' ? name : `${prefix}_${name}`;
            };
            for (const expression of expressions) {
                const type = expression.__class__.__name__;
                for (const input of expression.inputs) {
                    input._fullname = getName(context, getFullName(namePrefix, input._name));
                }
                for (const output of expression.outputs) {
                    output._fullname = getName(context, getFullName(namePrefix, output._name));
                }
                switch (type) {
                    case 'Input': {
                        break;
                    }
                    case 'GetAttr': {
                        if (expression.outputs[0].__class__.__name__ === 'TensorNode') {
                            const tensor = parseGetAttr(tmodule, expression);
                            const type = getTensorType(tensor.dtype, tensor.data.shape);
                            const data = tensor.data.data;
                            expression.outputs[0].initializer = new megengine.Tensor(expression.name, type, data);
                        }
                        break;
                    }
                    case 'Constant': {
                        if (expression.outputs[0].__class__.__name__ === 'TensorNode') {
                            const tensor = expression.value;
                            const type = getTensorType(tensor.dtype, tensor.data.shape);
                            const data = tensor.data.data;
                            expression.outputs[0].initializer = new megengine.Tensor('', type, data);
                        }
                        break;
                    }
                    case 'CallMethod': {
                        const obj = { 'name': '' };
                        let state = {};
                        if (expression.method === '__call__') {
                            const module = parseGetAttr(tmodule, expression.inputs[0].expr);
                            const getModuleType = (obj) => {
                                if (obj.module !== undefined) {
                                    return `${obj.module[0]}.${obj.module[1]}`;
                                }
                                return `${obj.__class__.__module__}.${obj.__class__.__name__}`;
                            };
                            const moduleType = module.__class__.__name__ === 'TracedModule' ? 'TracedModule' : getModuleType(module);
                            if (moduleType === 'TracedModule') {
                                const moduleName = expression.outputs[0]._name.endsWith("_out") ? expression.outputs[0]._name.substring(0, expression.outputs[0]._name.length - 4) : expression.outputs[0]._name;
                                const prefix = getFullName(namePrefix, moduleName);
                                const internalGraph = module.argdef_graph_map.get(expression.arg_def);
                                for (let i = 0; i < expression.inputs.length; i++) {
                                    const actualName = getFullName(namePrefix, expression.inputs[i]._name);
                                    const internalName = getFullName(prefix, internalGraph._inputs[i]._name);
                                    context.set(internalName, actualName);
                                }
                                for (let i = 0; i < expression.outputs.length; i++) {
                                    const actualName = getFullName(namePrefix, expression.outputs[i]._name);
                                    const internalName = getFullName(prefix, internalGraph._outputs[i]._name);
                                    if (context.get(internalName) === undefined) {
                                        context.set(internalName, actualName);
                                    } else {
                                        context.set(actualName, context.get(internalName));
                                    }
                                }
                                loadGraph(module, internalGraph, context, prefix, metadata, false);
                                continue;
                            }
                            obj.type = moduleType;
                            state = module.__class__.__name__ === 'TracedModule' ? module : module.state;
                            if (state === undefined) {
                                state = module;
                            }
                        } else {
                            obj.type = expression.method;
                            const [args, kwargs] = expression.arg_def.children_defs;
                            const schema = metadata.type(expression.method);
                            state = parseArgs(args, kwargs, schema);
                        }
                        const node = new megengine.Node(metadata, obj, values, null, expression, state);
                        this.nodes.push(node);
                        break;
                    }
                    case 'CallFunction': {
                        const getFunctionType = (obj) => {
                            if (obj.func.__module__ !== undefined) {
                                return `${obj.func.__module__}.${obj.func.__name__}`;
                            }
                            return `${obj.func[0]}.${obj.func[1]}`;
                        };
                        const func = getFunctionType(expression);
                        const item = { 'name': '', 'type': func };
                        const [args, kwargs] = expression.arg_def.children_defs;
                        const schema = metadata.type(func);
                        const state = parseArgs(args, kwargs, schema);
                        const node = new megengine.Node(metadata, item, values, null, expression, state);
                        this.nodes.push(node);
                        break;
                    }
                    case 'Apply': {
                        const opdef = expression.opdef_state ? expression.opdef_state.get('opdef_type') : expression.opdef.type;
                        const item = { 'name': '', 'type': `${opdef.__module__}.${opdef.__name__}` };
                        const node = new megengine.Node(metadata, item, values, null, expression, expression.opdef_state);
                        this.nodes.push(node);
                        break;
                    }
                    default: {
                        break;
                    }
                }
            }
        };
        if (obj.argdef_graph_map) {
            const [graph] = Array.from(obj.argdef_graph_map.values());
            loadGraph(obj, graph, new Map(), '', metadata, true);
            return;
        }
        const extraInfoNameset = new Set();
        const getExtraInfo = (opr) => {
            let name = opr.name;
            let repeatIdx = 0;
            while (extraInfoNameset.has(name)) {
                for (const id of opr.inputs) {
                    name = `${name}[${id}]`;
                }
                name += repeatIdx;
                repeatIdx += 1;
            }
            extraInfoNameset.add(name);
            const type = opr.type.replace(/V(\d+)$/, '');
            const args = [];
            if (opr.tensors.length > 0) {
                const [tensor] = opr.tensors;
                const type = new megengine.TensorType(tensor.dtype.type, new megengine.TensorShape(tensor.shape));
                const data = tensor.data.byteLength === 0 ? undefined : tensor.data.slice(0);
                const initializer = opr.type === 'Host2DeviceCopy' ? undefined : new megengine.Tensor('', type, data);
                const quantization = tensor.dtype.param ? { scale: tensor.dtype.param.scale, zeroPoint: tensor.dtype.param.zero_point } : null;
                const value = values.map(name, type, initializer, quantization);
                args.push(value);
            } else {
                const type = opr.shape ? new megengine.TensorType('?', new megengine.TensorShape(opr.shape)) : null;
                const value = values.map(name, type);
                args.push(value);
            }
            return { name, type, args };
        };
        const getAllOprAndTensor = (oprs) => {
            const allOprAndTensor = new Map();
            for (const opr of oprs) {
                if (opr.type === 'MultipleDeviceTensorWithFormatHolder' || opr.outputs.length > 1) {
                    if (opr.type === 'MultipleDeviceTensorWithFormatHolder' || opr.type === 'MultipleDeviceTensorHolder') {
                        opr.type = 'ImmutableTensor';
                    }
                    for (let id = 0; id < opr.outputs.length; id++) {
                        const keyId = opr.outputs[id];
                        const name = obj.middle_tensors[keyId] ? obj.middle_tensors[keyId].name : String(keyId);
                        const type = opr.type;
                        const tensors = opr.tensors.length ? [opr.tensors[id]] : [];
                        const onlyShape = obj.middle_tensors[keyId] ? obj.middle_tensors[keyId].shape : [];
                        allOprAndTensor.set(keyId, { name, type, tensors, shape: onlyShape, inputs: opr.inputs, outputs: opr.outputs });
                        const _opr = allOprAndTensor.get(keyId);
                        _opr.extraInfo = getExtraInfo(_opr);
                    }
                } else {
                    const [keyId] = opr.outputs;
                    opr.name = obj.middle_tensors[keyId] ? obj.middle_tensors[keyId].name : String(keyId);
                    if (obj.middle_tensors[keyId] && obj.middle_tensors[keyId].shape) {
                        opr.shape = obj.middle_tensors[keyId].shape;
                    }
                    allOprAndTensor.set(keyId, opr);
                    const _opr = allOprAndTensor.get(keyId);
                    _opr.extraInfo = getExtraInfo(_opr);
                }
            }
            return allOprAndTensor;
        };
        const allOprAndTensor = getAllOprAndTensor(obj.oprs);
        for (const op of Array.from(allOprAndTensor.values())) {
            if (op.type === 'Host2DeviceCopy') {
                const argument = new megengine.Argument('input', op.extraInfo.args);
                this.inputs.push(argument);
            } else if (op.type !== 'ImmutableTensor') {
                this.nodes.push(new megengine.Node(metadata, op, values, allOprAndTensor));
            }
        }
        for (let i = 0; i < obj.output_vars_idx.length; i++) {
            const id = obj.output_vars_idx[i].compact_id;
            const out_type = `output${i === 0 ? '' : i}`;
            const argument = new megengine.Argument(out_type, allOprAndTensor.get(id).extraInfo.args);
            this.outputs.push(argument);
        }
    }
};

megengine.Argument = class {

    constructor(name, value, type = null, visible = true) {
        this.name = name;
        this.value = value;
        this.type = type;
        this.visible = visible;
    }
};

megengine.Value = class {

    constructor(name, type, initializer, quantization) {
        if (typeof name !== 'string') {
            throw new megengine.Error(`Invalid value identifier '${JSON.stringify(name)}'.`);
        }
        this.name = name;
        this.type = !type && initializer ? initializer.type : type;
        this.initializer = initializer;
        if (quantization && ((quantization.scale !== undefined && quantization.scale !== 1) || (quantization.zeroPoint !== undefined && quantization.zeroPoint !== 0))) {
            this.quantization = {
                type: 'linear',
                scale: [quantization.scale],
                offset: [quantization.zeroPoint]
            };
        }
    }
};

megengine.Node = class {

    constructor(metadata, obj, values, allOprAndTensor, expr, state) {
        this.name = '';
        const type = metadata.type(obj.type);
        this.type = type ? { ...type } : { name: obj.type };
        this.type.name = this.type.name.replace(/V(\d+)$/, '');
        if (this.type.name.length > 4 && this.type.name.startsWith('__') && this.type.name.endsWith('__')) {
            this.type.name = this.type.name.substring(2, this.type.name.length - 2);
        }
        this.inputs = [];
        this.outputs = [];
        this.chain = [];
        this.attributes = [];
        const attributes = [];
        if (obj.inputs && obj.outputs) {
            const inputSchemas = this.type && this.type.inputs ? [...this.type.inputs] : [];
            for (let i = 0; i < obj.inputs.length; i++) {
                const inputOpr = allOprAndTensor.get(obj.inputs[i]);
                const schema = inputSchemas.length > 0 ? inputSchemas.shift() : { name: `input${i === 0 ? '' : i.toString()}` };
                const argument = new megengine.Argument(schema.name, inputOpr.extraInfo.args);
                this.inputs.push(argument);
            }
            const outputSchemas = this.type && this.type.outputs ? [...this.type.outputs] : [];
            for (let i = 0; i < obj.outputs.length; i++) {
                const outputOpr = allOprAndTensor.get(obj.outputs[i]);
                const schema = outputSchemas.length > 0 ? outputSchemas.shift() : { name: `output${i === 0 ? '' : i.toString()}` };
                const argument = new megengine.Argument(schema.name, outputOpr.extraInfo.args);
                this.outputs.push(argument);
            }
            if (obj.param) {
                for (const [name, value] of Object.entries(obj.param)) {
                    if (value !== null) {
                        const schema = metadata.attribute(obj.param.constructor.name, name);
                        attributes.push([schema, name, value]);
                    }
                }
            }
        }
        if (expr) {
            let inpIdx = 0;
            for (const i of expr.inputs) {
                if (i.__class__.__name__ !== 'ModuleNode') {
                    const initializer = i.initializer === undefined ? null : i.initializer;
                    const name = `input${inpIdx === 0 ? '' : inpIdx}`;
                    const dtype = i._dtype ? i._dtype.__name__ : null;
                    const shape = new megengine.TensorShape(i._shape);
                    const type = new megengine.TensorType(dtype, shape);
                    const value = values.map(i._fullname, type, initializer);
                    const argument = new megengine.Argument(name, [value]);
                    this.inputs.push(argument);
                    inpIdx += 1;
                }
            }
            let outIdx = 0;
            let qparams = null;
            for (const o of expr.outputs) {
                if (o._qparams !== null) {
                    qparams = o._qparams[1];
                }
                const name = `output${outIdx === 0 ? '' : outIdx}`;
                const dtype = o._dtype ? o._dtype.__name__ : null;
                const shape = new megengine.TensorShape(o._shape);
                const type = new megengine.TensorType(dtype, shape);
                const value = values.map(o._fullname, type, null);
                const argument = new megengine.Argument(name, [value]);
                this.outputs.push(argument);
                outIdx += 1;
            }
            if (qparams !== null) {
                state = state === null ? {} : state;
                state.scale = qparams.scale;
                state.zero_point = qparams.zero_point;
                state.quant_dtype_meta = qparams.dtype_meta;
            }
            if (state !== null) {
                for (const [key, obj] of Array.from(state)) {
                    const isModule = (obj) => {
                        return obj && (obj.state || obj._forward_pre_hooks);
                    };
                    const isTensor = (obj) => {
                        return obj && obj.__class__ && obj.__class__.__module__ === 'megengine.tensor' && (obj.__class__.__name__ === 'Tensor' || obj.__class__.__name__ === 'Parameter');
                    };
                    if (!key.startsWith('_') && !isModule(obj)) {
                        if (isTensor(obj)) {
                            const tensor = obj;
                            const dtype = tensor.dtype ? tensor.dtype.__name__ : null;
                            const shape = new megengine.TensorShape(tensor.data.shape);
                            const type = new megengine.TensorType(dtype, shape);
                            const data = tensor.data.data;
                            const initializer = new megengine.Tensor(key, type, data);
                            const value = values.map('', type, initializer);
                            const argument = new megengine.Argument(key, [value]);
                            this.inputs.push(argument);
                        } else {
                            const value = obj === null ? 'None' : obj;
                            attributes.push([null, key, value]);
                        }
                    }
                }
            }
        }
        this.attributes = attributes.map(([metadata, name, value]) => {
            value = ArrayBuffer.isView(value) ? Array.from(value) : value;
            let type = metadata ? metadata.type : null;
            let visible = true;
            if (name === 'training') {
                visible = false;
                type = 'boolean';
            }
            if (megengine.schema) {
                if (megengine.schema.param[type]) {
                    const enumType = megengine.schema.param[type];
                    if (enumType) {
                        value = enumType[value] || value;
                    }
                } else if (megengine.schema[type]) {
                    const enumType = megengine.schema[type];
                    if (enumType) {
                        value = enumType[value] || value;
                    }
                } else if (megengine.schema.v2[type]) {
                    const enumType = megengine.schema.v2[type];
                    if (enumType) {
                        value = enumType[value] || value;
                    }
                }
            }
            return new megengine.Argument(name, value, type, visible);
        });
    }
};

megengine.Tensor = class {

    constructor(name, type, data) {
        this.category = 'Tensor';
        this.name = name || '';
        this.type = type;
        this.values = data;
    }
};

megengine.TensorType = class {

    constructor(dataType, shape) {
        if (typeof dataType === 'number') {
            dataType = megengine.schema.DTypeEnum[dataType] || dataType;
        }
        dataType = typeof dataType === 'string' ? dataType.toLowerCase() : dataType;
        megengine.TensorType._dataTypes = megengine.TensorType._dataTypes || new Map([
            ['bool', 'boolean'],
            ['byte', 'uint8'], ['quantizeds4asymm', 'uint8'], ['quantizeds8asymm', 'uint8'], ['uintb4', 'uint8'],
            ['quantizeds1', 'int8'], ['quantizeds4', 'int8'], ['quantizeds8', 'int8'], ['intb1', 'int8'], ['intb2', 'int8'], ['intb4', 'int8'], ['qint8', 'int8'],
            ['quantizeds16', 'int16'],
            ['quantizeds32', 'int32']
        ]);
        this.dataType = megengine.TensorType._dataTypes.get(dataType) || dataType;
        this.shape = shape;
    }

    equals(obj) {
        return obj && this.dataType === obj.dataType && this.shape && this.shape.equals(obj.shape);
    }

    toString() {
        return this.dataType + this.shape.toString();
    }
};

megengine.TensorShape = class {

    constructor(dimensions) {
        this.dimensions = Array.from(dimensions || []);
    }

    equals(obj) {
        return obj && Array.isArray(obj.dimensions) &&
            Array.isArray(this.dimensions) && this.dimensions.length === obj.dimensions.length
            && obj.dimensions.every((value, index) => this.dimensions[index] === value);
    }

    toString() {
        if (this.dimensions && this.dimensions.length > 0) {
            return `[${this.dimensions.map((dimension) => dimension.toString()).join(',')}]`;
        }
        return '';
    }
};

megengine.Error = class extends Error {

    constructor(message) {
        super(message);
        this.name = 'Error loading MegEngine model.';
    }
};

export const ModelFactory = megengine.ModelFactory;