import type { Bytecode, Constant, Instruction } from "./bytecode"
import type { ExceptionHandler } from "./exception"
import { type Frame } from "./frame"
import { OpCode } from "./opcode"
import { Scope } from "./scope"
import { type Value, toValue, toNumber, isTrue, isEqual, toString, wrapNative, isWrapped } from "./value"

type NativeFunction = (...args: Value[]) => Promise<Value> | Value

export class VM {
  pc = 0
  stopped = false
  stack: Value[] = []
  callStack: Frame[] = []
  exceptionHandlers: ExceptionHandler[] = []
  scope: Scope
  constants: Constant[] = []
  instructions: Instruction[] = []
  labels: Map<number, string> = new Map()
  nativeFunctions: Map<string, NativeFunction> = new Map()

  constructor(bytecode: Bytecode) {
    this.instructions = bytecode.instructions
    this.constants = bytecode.constants
    this.labels = bytecode.labels || new Map()
    this.scope = new Scope()
  }

  registerFunction(name: string, fn: NativeFunction | Function) {
    // If it's already a NativeFunction, use it directly
    // Otherwise, assume it's a JS function and wrap it
    const wrapped = isWrapped(fn) ? fn as NativeFunction : wrapNative(fn)
    this.nativeFunctions.set(name, wrapped)
  }

  registerValueFunction(name: string, fn: NativeFunction) {
    this.nativeFunctions.set(name, fn)
  }

  async run(): Promise<Value> {
    this.pc = 0
    this.stopped = false

    while (!this.stopped && this.pc < this.instructions.length) {
      const instruction = this.instructions[this.pc]!
      await this.execute(instruction)
      this.pc++
    }

    return this.stack[this.stack.length - 1] || toValue(null)
  }

  async execute(instruction: Instruction) /* throws */ {
    switch (instruction.op) {
      case OpCode.PUSH:
        const constIdx = instruction.operand as number
        const constant = this.constants[constIdx]

        if (!constant || constant.type === 'function_def')
          throw new Error(`Invalid constant index: ${constIdx}`)

        this.stack.push(constant)
        break

      case OpCode.POP:
        this.stack.pop()
        break

      case OpCode.DUP:
        this.stack.push(this.stack[this.stack.length - 1]!)
        break

      case OpCode.ADD:
        this.binaryOp((a, b) => toNumber(a) + toNumber(b))
        break

      case OpCode.SUB:
        this.binaryOp((a, b) => toNumber(a) - toNumber(b))
        break

      case OpCode.MUL:
        this.binaryOp((a, b) => toNumber(a) * toNumber(b))
        break

      case OpCode.DIV:
        this.binaryOp((a, b) => toNumber(a) / toNumber(b))
        break

      case OpCode.MOD:
        this.binaryOp((a, b) => toNumber(a) % toNumber(b))
        break

      case OpCode.EQ:
        this.comparisonOp((a, b) => isEqual(a, b))
        break

      case OpCode.NEQ:
        this.comparisonOp((a, b) => !isEqual(a, b))
        break

      case OpCode.LT:
        this.comparisonOp((a, b) => toNumber(a) < toNumber(b))
        break

      case OpCode.GT:
        this.comparisonOp((a, b) => toNumber(a) > toNumber(b))
        break

      case OpCode.LTE:
        this.comparisonOp((a, b) => toNumber(a) <= toNumber(b))
        break

      case OpCode.GTE:
        this.comparisonOp((a, b) => toNumber(a) >= toNumber(b))
        break

      case OpCode.NOT:
        const val = this.stack.pop()!
        this.stack.push({ type: 'boolean', value: !isTrue(val) })
        break

      case OpCode.HALT:
        this.stopped = true
        break

      case OpCode.LOAD: {
        const varName = instruction.operand as string
        const value = this.scope.get(varName)

        if (value === undefined)
          throw new Error(`Undefined variable: ${varName}`)

        this.stack.push(value)
        break
      }

      case OpCode.TRY_LOAD: {
        const varName = instruction.operand as string
        const value = this.scope.get(varName)

        if (value === undefined)
          this.stack.push(toValue(varName))
        else
          this.stack.push(value)

        break
      }

      case OpCode.STORE:
        const name = instruction.operand as string
        const toStore = this.stack.pop()!
        this.scope.set(name, toStore)
        break

      case OpCode.JUMP:
        this.pc += (instruction.operand as number)
        break

      case OpCode.JUMP_IF_FALSE:
        const cond = this.stack.pop()!
        if (!isTrue(cond))
          this.pc += (instruction.operand as number)
        break

      case OpCode.JUMP_IF_TRUE:
        const condTrue = this.stack.pop()!
        if (isTrue(condTrue))
          this.pc += (instruction.operand as number)
        break

      case OpCode.BREAK:
        // Unwind call stack until we find a frame marked as break target
        let foundBreakTarget = false
        while (this.callStack.length) {
          const frame = this.callStack.pop()!
          this.scope = frame.returnScope
          this.pc = frame.returnAddress
          if (frame.isBreakTarget) {
            foundBreakTarget = true
            break
          }
        }
        if (!foundBreakTarget)
          throw new Error('BREAK: no break target found')
        break

      case OpCode.PUSH_TRY:
        const catchAddress = instruction.operand as number
        this.exceptionHandlers.push({
          catchAddress,
          callStackDepth: this.callStack.length,
          scope: this.scope
        })
        break

      case OpCode.PUSH_FINALLY:
        const finallyAddress = instruction.operand as number
        if (this.exceptionHandlers.length === 0)
          throw new Error('PUSH_FINALLY: no exception handler to modify')

        this.exceptionHandlers[this.exceptionHandlers.length - 1]!.finallyAddress = finallyAddress
        break

      case OpCode.POP_TRY:
        if (this.exceptionHandlers.length === 0)
          throw new Error('POP_TRY: no exception handler to pop')

        this.exceptionHandlers.pop()
        break

      case OpCode.THROW:
        const errorValue = this.stack.pop()!

        if (this.exceptionHandlers.length === 0) {
          const errorMsg = toString(errorValue)
          throw new Error(`Uncaught exception: ${errorMsg}`)
        }

        const throwHandler = this.exceptionHandlers.pop()!

        while (this.callStack.length > throwHandler.callStackDepth)
          this.callStack.pop()

        this.scope = throwHandler.scope

        this.stack.push(errorValue)

        // Jump to finally if present, otherwise jump to catch
        const targetAddress = throwHandler.finallyAddress !== undefined
          ? throwHandler.finallyAddress
          : throwHandler.catchAddress

        // subtract 1 because pc will be incremented
        this.pc = targetAddress - 1
        break

      case OpCode.MAKE_ARRAY:
        const arraySize = instruction.operand as number
        const items: Value[] = []

        for (let i = 0; i < arraySize; i++)
          items.unshift(this.stack.pop()!)

        this.stack.push({ type: 'array', value: items })
        break

      case OpCode.ARRAY_GET:
        const index = this.stack.pop()!
        const array = this.stack.pop()!

        if (array.type !== 'array')
          throw new Error('ARRAY_GET: not an array')

        const idx = Math.floor(toNumber(index))
        if (idx < 0 || idx >= array.value.length)
          throw new Error(`ARRAY_GET: index ${idx} out of bounds`)

        this.stack.push(array.value[idx]!)
        break

      case OpCode.ARRAY_SET:
        const setValue = this.stack.pop()!
        const setIndex = this.stack.pop()!
        const setArray = this.stack.pop()!

        if (setArray.type !== 'array')
          throw new Error('ARRAY_SET: not an array')

        const setIdx = Math.floor(toNumber(setIndex))
        if (setIdx < 0 || setIdx >= setArray.value.length)
          throw new Error(`ARRAY_SET: index ${setIdx} out of bounds`)

        setArray.value[setIdx] = setValue
        break

      case OpCode.ARRAY_PUSH:
        const pushValue = this.stack.pop()!
        const pushArray = this.stack.pop()!

        if (pushArray.type !== 'array')
          throw new Error('ARRAY_PUSH: not an array')

        pushArray.value.push(pushValue)
        break

      case OpCode.ARRAY_LEN:
        const lenArray = this.stack.pop()!

        if (lenArray.type !== 'array')
          throw new Error('ARRAY_LEN: not an array')

        this.stack.push({ type: 'number', value: lenArray.value.length })
        break

      case OpCode.MAKE_DICT:
        const dictPairs = instruction.operand as number
        const dict = new Map<string, Value>()

        for (let i = 0; i < dictPairs; i++) {
          const value = this.stack.pop()!
          const key = this.stack.pop()!
          dict.set(toString(key), value)
        }

        this.stack.push({ type: 'dict', value: dict })
        break

      case OpCode.DICT_GET:
        const getKey = this.stack.pop()!
        const getDict = this.stack.pop()!

        if (getDict.type !== 'dict')
          throw new Error('DICT_GET: not a dict')

        this.stack.push(getDict.value.get(toString(getKey)) || toValue(null))
        break

      case OpCode.DICT_SET:
        const dictSetValue = this.stack.pop()!
        const dictSetKey = this.stack.pop()!
        const dictSet = this.stack.pop()!

        if (dictSet.type !== 'dict')
          throw new Error('DICT_SET: not a dict')

        dictSet.value.set(toString(dictSetKey), dictSetValue)
        break

      case OpCode.DICT_HAS:
        const hasKey = this.stack.pop()!
        const hasDict = this.stack.pop()!

        if (hasDict.type !== 'dict')
          throw new Error('DICT_HAS: not a dict')

        this.stack.push({ type: 'boolean', value: hasDict.value.has(toString(hasKey)) })
        break

      case OpCode.MAKE_FUNCTION:
        const fnDefIdx = instruction.operand as number
        const fnDef = this.constants[fnDefIdx]

        if (!fnDef || fnDef.type !== 'function_def')
          throw new Error(`Invalid function definition index: ${fnDefIdx}`)

        this.stack.push({
          type: 'function',
          params: fnDef.params,
          defaults: fnDef.defaults,
          body: fnDef.body,
          variadic: fnDef.variadic,
          named: fnDef.named,
          parentScope: this.scope,
          value: '<function>'
        })
        break

      case OpCode.CALL: {
        // Pop named count from stack (top)
        const namedCount = toNumber(this.stack.pop()!)

        // Pop positional count from stack
        const positionalCount = toNumber(this.stack.pop()!)

        // Pop named arguments (name-value pairs) from stack
        // Stack has: ... key1 value1 key2 value2 (top)
        // So we pop value2, key2, value1, key1
        const namedArgs = new Map<string, Value>()
        const namedPairs: Array<{ key: string; value: Value }> = []
        for (let i = 0; i < namedCount; i++) {
          const value = this.stack.pop()!
          const key = this.stack.pop()!
          namedPairs.unshift({ key: toString(key), value })
        }
        for (const pair of namedPairs) {
          namedArgs.set(pair.key, pair.value)
        }

        // Pop positional arguments from stack
        const positionalArgs: Value[] = []
        for (let i = 0; i < positionalCount; i++)
          positionalArgs.unshift(this.stack.pop()!)

        const fn = this.stack.pop()!

        if (fn.type !== 'function')
          throw new Error('CALL: not a function')

        if (this.callStack.length > 0)
          this.callStack[this.callStack.length - 1]!.isBreakTarget = true

        this.callStack.push({
          returnAddress: this.pc,
          returnScope: this.scope,
          isBreakTarget: false
        })

        this.scope = new Scope(fn.parentScope)

        // Determine how many params are fixed (excluding variadic and named)
        let fixedParamCount = fn.params.length
        if (fn.variadic) fixedParamCount--
        if (fn.named) fixedParamCount--

        // Bind fixed parameters using priority: named arg > positional arg > default > null
        for (let i = 0; i < fixedParamCount; i++) {
          const paramName = fn.params[i]!

          // Check if named argument was provided for this param
          if (namedArgs.has(paramName)) {
            this.scope.set(paramName, namedArgs.get(paramName)!)
            namedArgs.delete(paramName) // Remove from named args so it won't go to named
          } else if (positionalArgs[i] !== undefined) {
            this.scope.set(paramName, positionalArgs[i]!)
          } else if (fn.defaults[paramName] !== undefined) {
            const defaultIdx = fn.defaults[paramName]!
            const defaultValue = this.constants[defaultIdx]!
            if (defaultValue.type === 'function_def')
              throw new Error('Default value cannot be a function definition')
            this.scope.set(paramName, defaultValue)
          } else {
            this.scope.set(paramName, toValue(null))
          }
        }

        // Handle variadic parameter (collect remaining positional args)
        if (fn.variadic) {
          const variadicParamName = fn.params[fn.params.length - (fn.named ? 2 : 1)]!
          const remainingArgs = positionalArgs.slice(fixedParamCount)
          this.scope.set(variadicParamName, { type: 'array', value: remainingArgs })
        }

        // Handle named parameter (collect remaining named args that didn't match params)
        if (fn.named) {
          const namedParamName = fn.params[fn.params.length - 1]!
          const namedDict = new Map<string, Value>()
          for (const [key, value] of namedArgs) {
            namedDict.set(key, value)
          }
          this.scope.set(namedParamName, { type: 'dict', value: namedDict })
        }

        // subtract 1 because pc was incremented
        this.pc = fn.body - 1
        break
      }

      case OpCode.TAIL_CALL: {
        // Pop named count from stack (top)
        const tailNamedCount = toNumber(this.stack.pop()!)

        // Pop positional count from stack
        const tailPositionalCount = toNumber(this.stack.pop()!)

        // Pop named arguments (name-value pairs) from stack
        const tailNamedArgs = new Map<string, Value>()
        const tailNamedPairs: Array<{ key: string; value: Value }> = []
        for (let i = 0; i < tailNamedCount; i++) {
          const value = this.stack.pop()!
          const key = this.stack.pop()!
          tailNamedPairs.unshift({ key: toString(key), value })
        }
        for (const pair of tailNamedPairs) {
          tailNamedArgs.set(pair.key, pair.value)
        }

        // Pop positional arguments from stack
        const tailPositionalArgs: Value[] = []
        for (let i = 0; i < tailPositionalCount; i++)
          tailPositionalArgs.unshift(this.stack.pop()!)

        const tailFn = this.stack.pop()!

        if (tailFn.type !== 'function')
          throw new Error('TAIL_CALL: not a function')

        this.scope = new Scope(tailFn.parentScope)

        // Determine how many params are fixed (excluding variadic and named)
        let tailFixedParamCount = tailFn.params.length
        if (tailFn.variadic) tailFixedParamCount--
        if (tailFn.named) tailFixedParamCount--

        // Bind fixed parameters
        for (let i = 0; i < tailFixedParamCount; i++) {
          const paramName = tailFn.params[i]!

          if (tailNamedArgs.has(paramName)) {
            this.scope.set(paramName, tailNamedArgs.get(paramName)!)
            tailNamedArgs.delete(paramName)
          } else if (tailPositionalArgs[i] !== undefined) {
            this.scope.set(paramName, tailPositionalArgs[i]!)
          } else if (tailFn.defaults[paramName] !== undefined) {
            const defaultIdx = tailFn.defaults[paramName]!
            const defaultValue = this.constants[defaultIdx]!
            if (defaultValue.type === 'function_def')
              throw new Error('Default value cannot be a function definition')
            this.scope.set(paramName, defaultValue)
          } else {
            this.scope.set(paramName, toValue(null))
          }
        }

        // Handle variadic parameter
        if (tailFn.variadic) {
          const variadicParamName = tailFn.params[tailFn.params.length - (tailFn.named ? 2 : 1)]!
          const remainingArgs = tailPositionalArgs.slice(tailFixedParamCount)
          this.scope.set(variadicParamName, { type: 'array', value: remainingArgs })
        }

        // Handle named parameter
        if (tailFn.named) {
          const namedParamName = tailFn.params[tailFn.params.length - 1]!
          const namedDict = new Map<string, Value>()
          for (const [key, value] of tailNamedArgs) {
            namedDict.set(key, value)
          }
          this.scope.set(namedParamName, { type: 'dict', value: namedDict })
        }

        // subtract 1 because PC was incremented
        this.pc = tailFn.body - 1
        break
      }

      case OpCode.RETURN:
        const returnValue = this.stack.length ? this.stack.pop()! : toValue(null)

        const frame = this.callStack.pop()
        if (!frame)
          throw new Error('RETURN: no call frame to return from')

        this.scope = frame.returnScope
        this.pc = frame.returnAddress

        this.stack.push(returnValue)
        break

      case OpCode.CALL_NATIVE:
        const functionName = instruction.operand as string
        const tsFunction = this.nativeFunctions.get(functionName)

        if (!tsFunction)
          throw new Error(`CALL_NATIVE: function not found: ${functionName}`)

        // Mark current frame as break target (like CALL does)
        if (this.callStack.length > 0)
          this.callStack[this.callStack.length - 1]!.isBreakTarget = true

        // Pop all arguments from stack (TypeScript function consumes entire stack)
        const tsArgs = [...this.stack]
        this.stack = []

        // Call the TypeScript function and await if necessary
        const tsResult = await tsFunction(...tsArgs)

        // Push result back onto stack
        this.stack.push(tsResult)
        break

      default:
        throw `Unknown op: ${instruction.op}`
    }
  }

  binaryOp(fn: (a: Value, b: Value) => number) {
    const b = this.stack.pop()!
    const a = this.stack.pop()!
    const result = fn(a, b)
    this.stack.push({ type: 'number', value: result })
  }

  comparisonOp(fn: (a: Value, b: Value) => boolean) {
    const b = this.stack.pop()!
    const a = this.stack.pop()!
    const result = fn(a, b)
    this.stack.push({ type: 'boolean', value: result })
  }
}