import {
    type ConstantNode,
    isAccessorNode,
    isArrayNode,
    isConstantNode,
    isObjectNode,
    type FunctionNode,
    type MathNode,
} from 'mathjs';
import type { Options, Result } from './interface.js';
import type { State } from './state.js';
import { constantValue, equalText, globalFnName, len, symbolName } from './utils.js';
import { migrateAtomic, migrateExpr, migrateNode } from './node.js';
import { concat } from './concat.js';
import { toBoolean, toNumber, toString } from './to-type.js';
import { isVmArray, isVmPrimitive } from '@mirascript/mirascript';

const EXACT_UNARY_FUNCTION = new Map<string, Result['type']>([
    // 只支持标量的 mathjs 函数
    ['sin', 'number'],
    ['cos', 'number'],
    ['tan', 'number'],
    ['sinh', 'number'],
    ['cosh', 'number'],
    ['tanh', 'number'],
    ['asin', 'number'],
    ['acos', 'number'],
    ['atan', 'number'],
    ['asinh', 'number'],
    ['acosh', 'number'],
    ['atanh', 'number'],
    ['sqrt', 'number'],
    ['cbrt', 'number'],
    ['log', 'number'],
    ['gamma', 'number'],

    // 自定义函数，在 mirascript 也提供
    ['keys', 'array'],
    ['values', 'array'],
    ['entries', 'array'],

    // XStudio 提供的函数
    ['service', 'extern'],
]);

/** 只接受一个参数，返回数字的函数 */
const ENTRYWISE_UNARY_MATH_FUNCTION = new Set<string>(['abs', 'sign', 'factorial']);
/** 舍入函数 */
const ROUNDING_FUNCTION = {
    round: 'round',
    ceil: 'ceil',
    floor: 'floor',
    fix: 'trunc',
};
/** 数字断言函数 */
const NUMBER_ASSERT_FUNCTION = {
    isNaN: (c: Result) => `${c.code} is nan`,
    isInteger: (c: Result) => `((${c.code}) % 1) == 0`,
    isFinite: (c: Result) => `${c.code} is > -inf and < inf`,
};
/** 数字计算函数 */
const NUMERIC_MAT_FUNCTION = {
    add: (s, a, b) => `${globalFnName(s, 'matrix')}.add(${a.code}, ${b.code})`,
    subtract: (s, a, b) => `${globalFnName(s, 'matrix')}.subtract(${a.code}, ${b.code})`,
    multiply: (s, a, b) => `${globalFnName(s, 'matrix')}.multiply(${a.code}, ${b.code})`,
    divide: (s, a, b) =>
        `${globalFnName(s, 'matrix')}.multiply(${a.code}, ${globalFnName(s, 'matrix')}.invert(${b.code}))`,
    dotMultiply: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise_multiply(${a.code}, ${b.code})`,
    dotDivide: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise_divide(${a.code}, ${b.code})`,
    compare: (s, a, b) =>
        `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { if a < b { -1 } else if a > b { 1 } else { 0 } })`,
} satisfies Record<string, (state: State, a: Result, b: Result) => string>;

/** 数字计算函数 */
const BOOLEAN_MAT_FUNCTION = {
    equal: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a =~ b })`,
    unequal: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a !~ b })`,
    smaller: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a < b })`,
    smallerEq: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a <= b })`,
    larger: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a > b })`,
    largerEq: (s, a, b) => `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { a >= b })`,
    and: (s, a, b) => {
        return `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { ${toBoolean(s, { code: 'a' }).code} && ${toBoolean(s, { code: 'b' }).code} })`;
    },
    or: (s, a, b) => {
        return `${globalFnName(s, 'matrix')}.entrywise(${a.code}, ${b.code}, fn (a, b) { ${toBoolean(s, { code: 'a' }).code} || ${toBoolean(s, { code: 'b' }).code} })`;
    },
} satisfies Record<string, (state: State, a: Result, b: Result) => string>;

/** 兜底 */
function call(state: State, node: FunctionNode<MathNode>, done: (fn: Result, args: readonly Result[]) => void): Result {
    const { fn, args } = node;
    const fnRet = migrateExpr(state, fn);
    const argList = args.map((a) => migrateAtomic(state, a));
    done(fnRet, argList);
    return {
        code: `${fnRet.code}(${argList.map((a) => a.code).join(', ')})`,
    };
}

/** 函数调用 */
function migrateFunctionCall(
    state: State,
    node: FunctionNode<MathNode>,
    fnName: string,
    args: readonly MathNode[],
    options: Options,
): Result {
    const openE = options.format !== 'no-paren' ? '(' : '';
    const closeE = options.format !== 'no-paren' ? ')' : '';
    const g = (fnName: string) => globalFnName(state, fnName);

    if (EXACT_UNARY_FUNCTION.has(fnName) && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        return {
            type: EXACT_UNARY_FUNCTION.get(fnName)!,
            code: `${fnName}(${arg.code})`,
        };
    } else if (fnName === 'random' && args.length === 0) {
        return {
            type: 'number',
            code: `${g('random')}()`,
        };
    } else if (ENTRYWISE_UNARY_MATH_FUNCTION.has(fnName) && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        if (arg.type === 'number' || arg.type === 'boolean' || arg.type === 'string') {
            return {
                type: 'number',
                code: `${g(fnName)}(${arg.code})`,
            };
        }
        return {
            type: arg.type,
            code: `${g('matrix')}.entrywise(${arg.code}, nil, ${g(fnName)})`,
        };
    } else if (fnName in ROUNDING_FUNCTION && (args.length === 1 || args.length === 2)) {
        const arg = migrateAtomic(state, args[0]!);
        const n = args.length === 2 ? migrateAtomic(state, args[1]!) : undefined;
        const fn = ROUNDING_FUNCTION[fnName as keyof typeof ROUNDING_FUNCTION];
        if (arg.type === 'number' || arg.type === 'boolean' || arg.type === 'string') {
            return {
                type: 'number',
                code: `${g(fn)}(${arg.code}${n ? `, ${n.code}` : ''})`,
            };
        }
        return {
            type: arg.type,
            code: `${g('matrix')}.entrywise(${arg.code}, ${n ? n.code : 'nil'}, ${g(fn)})`,
        };
    } else if (fnName === 'max' || fnName === 'min') {
        if (args.length === 1) {
            const arg = migrateAtomic(state, args[0]!);
            return {
                type: 'number',
                code: `${g(fnName)}(${arg.code})`,
            };
        } else {
            const argList = args.map((a) => migrateAtomic(state, a));
            if (!argList.every((a) => a.type === 'number' || a.type === 'boolean' || a.type === 'string')) {
                state.warn(`函数行为可能不一致: ${fnName}`);
            }
            return {
                type: 'number',
                code: `${g(fnName)}(${argList.map((a) => a.code).join(', ')})`,
            };
        }
    } else if (fnName === 're' || fnName === 'im') {
        const m = migrateAtomic(state, args[0]!);
        state.err(`不支持复数`);
        return {
            type: 'number',
            code: `/* ${fnName} */(${m.code})`,
        };
    } else if (fnName === 'date' && args.length === 1) {
        return {
            type: 'number',
            code: `${g('to_timestamp')}(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'size' && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        return {
            type: arg.type,
            code: `${g('matrix')}.size(${arg.code})`,
        };
    } else if (fnName === 'inv' && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        return {
            type: arg.type,
            code: `${g('matrix')}.invert(${arg.code})`,
        };
    } else if (fnName === 'count' && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        if (arg.type !== 'string') {
            state.warn(`函数行为可能不一致: count`);
        }
        return len(state, arg);
    } else if (fnName === 'transpose' && args.length === 1) {
        const arg = migrateAtomic(state, args[0]!);
        return {
            type: arg.type,
            code: `${g('matrix')}.transpose(${arg.code})`,
        };
    } else if (fnName === 'diag' && (args.length === 1 || args.length === 2)) {
        const a = args.map((a) => migrateAtomic(state, a).code).join(', ');
        return {
            type: 'array',
            code: `${g('matrix')}.diagonal(${a})`,
        };
    } else if (fnName === 'zeros' || fnName === 'ones' || fnName === 'identity') {
        return {
            type: 'array',
            code: `${g('matrix')}.${fnName}(${args.map((a) => migrateAtomic(state, a).code).join(', ')})`,
        };
    } else if (fnName === 'concat') {
        const results = args.map((a) => migrateAtomic(state, a));
        if (results.some((r) => r.type === 'string')) return concat(state, results);
        if (results.some((r) => r.type === 'array')) {
            state.warn(`矩阵连接时结果可能不一致`);
            return {
                type: 'array',
                code: `${g('flatten')}([` + results.map((r) => r.code).join(', ') + `])`,
            };
        }
    } else if (fnName === 'numeric' && args.length === 1) {
        return toNumber(state, args[0]!);
    } else if (fnName === 'flatten' && args.length === 1) {
        return {
            type: 'array',
            code: `${g('flatten')}(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'norm' && args.length === 1) {
        if (isArrayNode(args[0])) {
            return {
                type: 'number',
                code: `${g('hypot')}(${args[0].items.map((a) => migrateAtomic(state, a).code).join(', ')})`,
            };
        }
        return {
            type: 'number',
            code: `${g('hypot')}(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'equalText' && args.length === 2) {
        const p = equalText(state, '==', args[0]!, args[1]!);
        return {
            type: p.type,
            code: `${openE}${p.code}${closeE}`,
        };
    } else if (fnName === 'toJson' && args.length === 1) {
        return {
            type: 'string',
            code: `${g('to_json')}(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'fromJson' && args.length === 1) {
        state.helper("fn @@from_json(s) { if type(s) == 'string' { from_json(s, s) } else { s } }");
        return {
            code: `@@from_json(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'is' && args.length === 2 && isConstantNode(args[1])) {
        const t = String((args[1] as ConstantNode<string | number>).value);
        switch (t) {
            case 'number':
            case 'string':
            case 'boolean':
                return {
                    type: 'boolean',
                    code: `${openE}type(${migrateAtomic(state, args[0]!).code}) == '${t}'${closeE}`,
                };
            case 'null':
            case 'undefined':
                return {
                    type: 'boolean',
                    code: `${openE}type(${migrateAtomic(state, args[0]!).code}) == 'nil'${closeE}`,
                };
            case 'Array':
            case 'array':
                state.warn(`'type' 与 'is' 结果可能不同`);
                return {
                    type: 'boolean',
                    code: `${openE}type(${migrateAtomic(state, args[0]!).code}) == 'array'${closeE}`,
                };
            case 'Object':
            case 'object': {
                const arg0 = migrateAtomic(state, args[0]!);
                if (arg0.code.endsWith('.value.rt_state')) {
                    return {
                        type: 'boolean',
                        code: `${openE}${arg0.code} != nil${closeE}`,
                    };
                }
                state.warn(`'type' 与 'is' 结果可能不同`);
                return {
                    type: 'boolean',
                    code: `${openE}type(${arg0.code}) == 'record'${closeE}`,
                };
            }
            case 'Function':
            case 'function': {
                if (symbolName(args[0]!) === 'service' && !state.locals.has('service')) {
                    return {
                        type: 'boolean',
                        code: `${openE}type(service) == 'extern'${closeE}`,
                    };
                }
                state.warn(`'type' 与 'is' 结果可能不同`);
                return {
                    type: 'boolean',
                    code: `${openE}type(${migrateAtomic(state, args[0]!).code}) == 'function'${closeE}`,
                };
            }
        }
    } else if (fnName === 'typeOf' && args.length === 1) {
        state.warn(`'type' 与 'typeOf' 结果可能不同`);
        return {
            type: 'string',
            code: `type(${migrateAtomic(state, args[0]!).code})`,
        };
    } else if (fnName === 'format' && args.length === 1) {
        return toString(state, args[0]!);
    } else if (fnName === 'print' && args.length === 1) {
        state.loose();
        return toString(state, args[0]!);
    } else if (
        fnName === 'print' &&
        args.length === 2 &&
        isConstantNode(args[0]) &&
        (isObjectNode(args[1]) || isArrayNode(args[1]))
    ) {
        const template = String(args[0].value);
        const values = isObjectNode(args[1])
            ? args[1].properties
            : (args[1].items as unknown as Record<string, MathNode>);
        return {
            type: 'string',
            code:
                '`' +
                template.replaceAll(/`|\$([\w.]+)/g, (original, key: string) => {
                    if (original === '`') return '\\`';
                    const keys = key.split('.').map((part) => {
                        const nPart = Number.parseInt(part);
                        if (!Number.isNaN(nPart) && part.length > 0 && nPart >= 1) {
                            return nPart - 1;
                        } else {
                            return part;
                        }
                    });
                    const value = values[keys.shift()!];
                    if (value == null) return original.replaceAll('$', String.raw`\$`);
                    if (!keys.length) {
                        return `$(${migrateAtomic(state, value).code})`;
                    }
                    return `$(${migrateNode(state, value, { format: 'paren' }).code}.${keys.join('.')})`;
                }) +
                '`',
        };
    } else if (fnName === 'string' && args.length === 1) {
        const inner = migrateAtomic(state, args[0]!);
        return {
            type: 'string',
            code: `${g('to_string')}(${inner.code})`,
        };
    } else if (fnName === 'sum' || fnName === 'prod') {
        const newFnName = fnName === 'sum' ? 'sum' : 'product';
        if (args.length !== 1) {
            return {
                type: 'number',
                code: `${g(newFnName)}(${args.map((a) => migrateAtomic(state, a).code).join(', ')})`,
            };
        }
        const arg = migrateAtomic(state, args[0]!);
        if (arg.literal) {
            if (isVmPrimitive(arg.literal)) {
                return {
                    type: 'number',
                    code: `${g(newFnName)}(${arg.code})`,
                };
            }
            if (isVmArray(arg.literal) && arg.literal.every((v) => isVmPrimitive(v))) {
                return {
                    type: 'number',
                    code: `${g(newFnName)}(${arg.code})`,
                };
            }
        }
        if (arg.type === 'array' || !arg.type) {
            return {
                type: 'number',
                code: `${g(newFnName)}(${arg.code}::${g('flatten')}())`,
            };
        }
        return {
            type: 'number',
            code: `${g(newFnName)}(${arg.code})`,
        };
    } else if (fnName in NUMBER_ASSERT_FUNCTION && args.length === 1) {
        const f = NUMBER_ASSERT_FUNCTION[fnName as keyof typeof NUMBER_ASSERT_FUNCTION];
        const arg = migrateAtomic(state, args[0]!);
        if (arg.type === 'number' || arg.type === 'boolean' || arg.type === 'string') {
            return {
                type: 'boolean',
                code: `${openE}${f(toNumber(state, arg))}${closeE}`,
            };
        }
        return {
            type: arg.type,
            code: `${g('matrix')}.entrywise(${arg.code}, nil, fn { ${f({ code: 'to_number(it)' })} })`,
        };
    } else if ((fnName in NUMERIC_MAT_FUNCTION || fnName in BOOLEAN_MAT_FUNCTION) && args.length === 2) {
        const f =
            NUMERIC_MAT_FUNCTION[fnName as keyof typeof NUMERIC_MAT_FUNCTION] ??
            BOOLEAN_MAT_FUNCTION[fnName as keyof typeof BOOLEAN_MAT_FUNCTION];
        const a = migrateAtomic(state, args[0]!);
        const b = migrateAtomic(state, args[1]!);
        if (
            (a.type === 'number' || a.type === 'boolean' || a.type === 'string') &&
            (b.type === 'number' || b.type === 'boolean' || b.type === 'string')
        ) {
            return {
                type: fnName in BOOLEAN_MAT_FUNCTION ? 'boolean' : 'number',
                code: f(state, a, b),
            };
        }
        return {
            type: 'array',
            code: f(state, a, b),
        };
    }

    return call(state, node, (fn, args) => {
        const fun = state.global(fnName);
        if (typeof fun == 'function') {
            if (args.some((a) => a.type !== 'number' && a.type !== 'boolean' && a.type !== 'string')) {
                state.warn(`函数行为可能不一致: ${fnName}`);
            }
        } else {
            state.err(`不支持的函数: ${fnName}`);
        }
    });
}

/** 转换 AST */
export function migrateCall(state: State, node: FunctionNode<MathNode>, options: Options): Result {
    const { fn, args } = node;

    const fnName = symbolName(fn);
    if (fnName && !state.locals.has(fnName)) {
        return migrateFunctionCall(state, node, fnName, args, options);
    }

    if (
        isAccessorNode(fn) &&
        fn.index.dimensions.length === 1 &&
        isConstantNode(fn.index.dimensions[0]) &&
        typeof fn.index.dimensions[0].value == 'string'
    ) {
        const thisArg = migrateExpr(state, fn.object);
        const fnName = String(fn.index.dimensions[0].value);
        if (thisArg.type === 'extern') {
            const argList = args.map((a) => migrateAtomic(state, a).code);
            return {
                code: `${thisArg.code}.${fnName}(${argList.join(', ')})`,
            };
        }
        const g = (fnName: string) => globalFnName(state, fnName);
        if (fnName === 'toString' && args.length === 0) {
            return {
                type: 'string',
                code: `${thisArg.code}::${g('to_string')}()`,
            };
        } else if (fnName === 'includes' && args.length === 1) {
            if (thisArg.type === 'array') {
                return {
                    type: 'boolean',
                    code: `(${migrateExpr(state, args[0]!).code} in ${thisArg.code})`,
                };
            }
            if (thisArg.type === 'string') {
                return {
                    type: 'boolean',
                    code: `${thisArg.code}::${g('contains')}(${migrateAtomic(state, args[0]!).code})`,
                };
            }
            state.helper(
                `fn @@includes(it, el) { match it::type() { case 'string' { it::contains(el) } case 'array' { el in it } case _ { false } } }`,
            );
            return {
                type: 'boolean',
                code: `${thisArg.code}::@@includes(${migrateAtomic(state, args[0]!).code})`,
            };
        } else if ((fnName === 'map' || fnName === 'filter') && args.length === 1) {
            return {
                type: 'array',
                code: `${thisArg.code}::${g(fnName)}(${migrateAtomic(state, args[0]!).code})`,
            };
        } else if (fnName === 'push' && args.length === 1) {
            state.warn(`'push' 的行为可能不一致`);
            return {
                type: 'array',
                code: `{ let ret = ${migrateAtomic(state, args[0]!).code}; ${thisArg.code} = [..${thisArg.code}, ret]; ret }`,
            };
        } else if (fnName === 'pop' && args.length === 0) {
            state.warn(`'pop' 的行为可能不一致`);
            return {
                type: 'array',
                code: `{ let ret = ${thisArg.code}[-1]; ${thisArg.code} = ${thisArg.code}[0..<-1]; ret }`,
            };
        } else if (fnName === 'find' && args.length === 1) {
            return {
                code: `${thisArg.code}::${g('find')}(${migrateAtomic(state, args[0]!).code}).1`,
            };
        } else if (fnName === 'findIndex' && args.length === 1) {
            return {
                type: 'number',
                code: `(${thisArg.code}::${g('find')}(${migrateAtomic(state, args[0]!).code}).0 ?? -1)`,
            };
        } else if (fnName === 'flatMap' && args.length === 1) {
            return {
                type: 'array',
                code: `${thisArg.code}::${g('map')}(${migrateAtomic(state, args[0]!).code})::flatten()`,
            };
        } else if (fnName === 'reverse' && args.length === 0) {
            return {
                type: 'array',
                code: `${thisArg.code}::${g('reverse')}()`,
            };
        } else if ((fnName === 'replaceAll' || fnName === 'replace') && args.length === 2) {
            if (fnName === 'replace') {
                state.warn(`MiraScript 的 'replace' 函数会替换全部结果`);
            }
            return {
                type: 'string',
                code: `${thisArg.code}::${g('replace')}(${migrateAtomic(state, args[0]!).code}, ${migrateAtomic(state, args[1]!).code})`,
            };
        } else if (fnName === 'split' && args.length === 1) {
            return {
                type: 'array',
                code: `${thisArg.code}::${g('split')}(${migrateAtomic(state, args[0]!).code})`,
            };
        } else if (fnName === 'repeat' && args.length === 1) {
            return {
                type: 'string',
                code: `${thisArg.code}::${g('repeat')}(${migrateAtomic(state, args[0]!).code})::${g('join')}()`,
            };
        } else if (fnName === 'concat') {
            if (thisArg.type === 'string') {
                return concat(state, [thisArg, ...args]);
            }
            if (thisArg.type === 'array') {
                return {
                    type: 'array',
                    code:
                        `${g('flatten')}([` +
                        [fn.object, ...args].map((a) => migrateAtomic(state, a).code).join(', ') +
                        `])`,
                };
            }
            state.helper(
                `fn @@concat(it, ..args) { let arr = ${g('flatten')}([it,..args]); match it::type() { case 'string' { arr::${g('join')}() } case _ { arr } } }`,
            );
            return {
                code: `${thisArg.code}::@@concat(${args.map((a) => migrateAtomic(state, a).code).join(', ')})`,
            };
        } else if (fnName === 'toFixed' && args.length <= 1) {
            const digits =
                args.length === 0 ? 0 : (constantValue(args[0]!) ?? `$(${migrateAtomic(state, args[0]!).code})`);
            return {
                type: 'string',
                code: `${thisArg.code}::${g('format')}('.${digits}')`,
            };
        } else if (fnName === 'join' && args.length <= 1) {
            let sep = ',';
            if (args[0]) {
                if (isConstantNode(args[0])) {
                    sep = String(args[0].value).replaceAll('`', '\\`');
                } else {
                    sep = `$(${migrateAtomic(state, args[0]).code})`;
                }
            }
            if (isArrayNode(fn.object)) {
                let ret = '';
                for (let i = 0; i < fn.object.items.length; i++) {
                    const item = fn.object.items[i]!;
                    if (i > 0) ret += sep;
                    if (isConstantNode(item)) {
                        const value = String(item.value).replaceAll('`', '\\`');
                        ret += value;
                    } else {
                        ret += `$(${migrateAtomic(state, item).code})`;
                    }
                }
                return {
                    type: 'string',
                    code: '`' + ret + '`',
                };
            }
            return {
                type: 'string',
                code: `${thisArg.code}::${g('join')}(\`${sep}\`)`,
            };
        } else if (fnName === 'toLowerCase' && args.length === 0) {
            return {
                type: 'string',
                code: `${thisArg.code}::${g('to_lowercase')}()`,
            };
        } else if (fnName === 'toUpperCase' && args.length === 0) {
            return {
                type: 'string',
                code: `${thisArg.code}::${g('to_uppercase')}()`,
            };
        }

        state.err(`不支持的方法: ${fnName}`);
        const argList = args.map((a) => migrateAtomic(state, a).code);
        return { code: `${thisArg.code}.${fnName}(${argList.join(', ')})` };
    }

    return call(state, node, () => {
        state.loose();
    });
}