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 } 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[] = []
  nativeFunctions: Map<string, NativeFunction> = new Map()

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

  registerFunction(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.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.CONTINUE:
        // Search for frame with continueAddress (don't pop yet)
        let continueFrame: Frame | undefined
        let frameIndex = this.callStack.length - 1
        while (frameIndex >= 0) {
          const frame = this.callStack[frameIndex]!
          if (frame.continueAddress !== undefined) {
            continueFrame = frame
            break
          }
          frameIndex--
        }

        if (!continueFrame)
          throw new Error('CONTINUE: no continue target found')

        // Pop all frames above the continue target
        while (this.callStack.length > frameIndex + 1)
          this.callStack.pop()

        // Restore scope and jump to continue address
        this.scope = continueFrame.returnScope
        this.pc = continueFrame.continueAddress! - 1  // -1 because PC will be incremented
        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,
          kwargs: fnDef.kwargs,
          parentScope: this.scope
        })
        break

      case OpCode.CALL:
        const argCount = instruction.operand as number

        const args: Value[] = []
        for (let i = 0; i < argCount; i++)
          args.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,
          continueAddress: undefined
        })

        this.scope = new Scope(fn.parentScope)

        for (let i = 0; i < fn.params.length; i++) {
          const paramName = fn.params[i]!
          const argValue = args[i]

          if (argValue !== undefined) {
            this.scope.set(paramName, argValue)
          } 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))
          }
        }

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

      case OpCode.TAIL_CALL:
        const tailArgCount = instruction.operand as number

        const tailArgs: Value[] = []
        for (let i = 0; i < tailArgCount; i++)
          tailArgs.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)

        // same logic as CALL
        for (let i = 0; i < tailFn.params.length; i++) {
          const paramName = tailFn.params[i]!
          const argValue = tailArgs[i]

          if (argValue !== undefined) {
            this.scope.set(paramName, argValue)
          } 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))
          }
        }

        // 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 })
  }
}