import { glob, readFileSync } from 'fs'
import { VM, fromValue, type Bytecode } from 'reefvm'
import { type Tree } from '@lezer/common'
import { Compiler } from '#compiler/compiler'
import { parser } from '#parser/shrimp'
import { setGlobals } from '#parser/tokenizer'
import { globals as shrimpGlobals, colors } from '#prelude'

export { Compiler } from '#compiler/compiler'
export { parser } from '#parser/shrimp'
export { globals } from '#prelude'

export class Shrimp {
    vm: VM
    private globals?: Record<string, any>

    constructor(globals?: Record<string, any>) {
        const emptyBytecode = { instructions: [], constants: [], labels: new Map() }
        this.vm = new VM(emptyBytecode, Object.assign({}, shrimpGlobals, globals ?? {}))
        this.globals = globals
    }

    async run(code: string | Bytecode, locals?: Record<string, any>): Promise<any> {
        let bytecode

        if (typeof code === 'string') {
            const compiler = new Compiler(code, Object.keys(Object.assign({}, shrimpGlobals, this.globals ?? {}, locals ?? {})))
            bytecode = compiler.bytecode
        } else {
            bytecode = code
        }

        if (locals) this.vm.pushScope(locals)
        this.vm.appendBytecode(bytecode)
        await this.vm.continue()
        if (locals) this.vm.popScope()

        return this.vm.stack.length ? fromValue(this.vm.stack.at(-1)!) : null
    }

    get(name: string): any {
        const value = this.vm.scope.get(name)
        return value ? fromValue(value) : null
    }
}

export async function runFile(path: string, globals?: Record<string, any>): Promise<any> {
    const code = readFileSync(path, 'utf-8')
    return await runCode(code, globals)
}

export async function runCode(code: string, globals?: Record<string, any>): Promise<any> {
    return await runBytecode(compileCode(code, globals), globals)
}

export async function runBytecode(bytecode: Bytecode, globals?: Record<string, any>): Promise<any> {
    try {
        const vm = new VM(bytecode, Object.assign({}, shrimpGlobals, globals))
        await vm.run()
        return vm.stack.length ? fromValue(vm.stack[vm.stack.length - 1]!) : null
    } catch (error: any) {
        console.error(`${colors.red}Error:${colors.reset} ${error.message}`)
        process.exit(1)
    }
}

export function compileFile(path: string, globals?: Record<string, any>): Bytecode {
    const code = readFileSync(path, 'utf-8')
    return compileCode(code, globals)
}

export function compileCode(code: string, globals?: Record<string, any>): Bytecode {
    const globalNames = [...Object.keys(shrimpGlobals), ...(globals ? Object.keys(globals) : [])]
    const compiler = new Compiler(code, globalNames)
    return compiler.bytecode
}

export function parseFile(path: string, globals?: Record<string, any>): Tree {
    const code = readFileSync(path, 'utf-8')
    return parseCode(code, globals)
}

export function parseCode(code: string, globals?: Record<string, any>): Tree {
    const globalNames = [...Object.keys(shrimpGlobals), ...(globals ? Object.keys(globals) : [])]
    setGlobals(globalNames)
    return parser.parse(code)
}