package vm

import (
	"fmt"

	"git.sharkk.net/Sharkk/Mako/types"
)

// Scope represents a lexical scope
type Scope struct {
	Variables map[string]types.Value
}

type VM struct {
	constants []any
	globals   map[string]types.Value
	scopes    []Scope // Stack of local scopes
	stack     []types.Value
	sp        int // Stack pointer
}

func New() *VM {
	return &VM{
		globals: make(map[string]types.Value),
		scopes:  []Scope{}, // Initially no scopes
		stack:   make([]types.Value, 1024),
		sp:      0,
	}
}

func (vm *VM) Run(bytecode *types.Bytecode) {
	vm.constants = bytecode.Constants

	for ip := 0; ip < len(bytecode.Instructions); ip++ {
		instruction := bytecode.Instructions[ip]

		switch instruction.Opcode {
		case types.OpConstant:
			constIndex := instruction.Operand
			constant := vm.constants[constIndex]

			switch v := constant.(type) {
			case string:
				vm.push(types.NewString(v))
			case float64:
				vm.push(types.NewNumber(v))
			case bool:
				vm.push(types.NewBoolean(v))
			case nil:
				vm.push(types.NewNull())
			}

		case types.OpSetLocal:
			constIndex := instruction.Operand
			name := vm.constants[constIndex].(string)
			value := vm.pop()

			// Set in current scope if it exists
			if len(vm.scopes) > 0 {
				vm.scopes[len(vm.scopes)-1].Variables[name] = value
			} else {
				// No scope, set as global
				vm.globals[name] = value
			}

		case types.OpGetLocal:
			constIndex := instruction.Operand
			name := vm.constants[constIndex].(string)

			// Check local scopes from innermost to outermost
			found := false
			for i := len(vm.scopes) - 1; i >= 0; i-- {
				if val, ok := vm.scopes[i].Variables[name]; ok {
					vm.push(val)
					found = true
					break
				}
			}

			// If not found in locals, check globals
			if !found {
				if val, ok := vm.globals[name]; ok {
					vm.push(val)
				} else {
					vm.push(types.NewNull())
				}
			}

		case types.OpSetGlobal:
			constIndex := instruction.Operand
			name := vm.constants[constIndex].(string)
			value := vm.pop()
			vm.globals[name] = value

		case types.OpGetGlobal:
			constIndex := instruction.Operand
			name := vm.constants[constIndex].(string)
			if val, ok := vm.globals[name]; ok {
				vm.push(val)
			} else {
				vm.push(types.NewNull())
			}

		case types.OpEnterScope:
			// Push a new scope
			vm.scopes = append(vm.scopes, Scope{
				Variables: make(map[string]types.Value),
			})

		case types.OpExitScope:
			// Pop the current scope
			if len(vm.scopes) > 0 {
				vm.scopes = vm.scopes[:len(vm.scopes)-1]
			}

		case types.OpNewTable:
			vm.push(types.NewTableValue())

		case types.OpSetIndex:
			value := vm.pop()
			key := vm.pop()
			tableVal := vm.pop()

			if tableVal.Type != types.TypeTable {
				fmt.Println("Error: attempt to index non-table value")
				vm.push(types.NewNull())
				continue
			}

			table := tableVal.Data.(*types.Table)
			table.Set(key, value)
			vm.push(tableVal)

		case types.OpGetIndex:
			key := vm.pop()
			tableVal := vm.pop()

			if tableVal.Type != types.TypeTable {
				fmt.Println("Error: attempt to index non-table value")
				vm.push(types.NewNull())
				continue
			}

			table := tableVal.Data.(*types.Table)
			value := table.Get(key)
			vm.push(value)

		case types.OpDup:
			if vm.sp > 0 {
				vm.push(vm.stack[vm.sp-1])
			}

		case types.OpPop:
			vm.pop()

		case types.OpEcho:
			value := vm.pop()
			switch value.Type {
			case types.TypeString:
				fmt.Println(value.Data.(string))
			case types.TypeNumber:
				fmt.Println(value.Data.(float64))
			case types.TypeBoolean:
				fmt.Println(value.Data.(bool))
			case types.TypeNull:
				fmt.Println("null")
			case types.TypeTable:
				fmt.Println(vm.formatTable(value.Data.(*types.Table)))
			}

		// Jump instructions
		case types.OpJumpIfFalse:
			condition := vm.pop()
			// Consider falsy: false, null, 0
			shouldJump := false

			if condition.Type == types.TypeBoolean && !condition.Data.(bool) {
				shouldJump = true
			} else if condition.Type == types.TypeNull {
				shouldJump = true
			} else if condition.Type == types.TypeNumber && condition.Data.(float64) == 0 {
				shouldJump = true
			}

			if shouldJump {
				ip = instruction.Operand - 1 // -1 because loop will increment
			}

		case types.OpJump:
			ip = instruction.Operand - 1 // -1 because loop will increment

		// Arithmetic operations
		case types.OpAdd:
			right := vm.pop()
			left := vm.pop()

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				result := left.Data.(float64) + right.Data.(float64)
				vm.push(types.NewNumber(result))
			} else if left.Type == types.TypeString && right.Type == types.TypeString {
				// String concatenation
				result := left.Data.(string) + right.Data.(string)
				vm.push(types.NewString(result))
			} else {
				fmt.Println("Error: cannot add values of different types")
				vm.push(types.NewNull())
			}

		case types.OpSubtract:
			right := vm.pop()
			left := vm.pop()

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				result := left.Data.(float64) - right.Data.(float64)
				vm.push(types.NewNumber(result))
			} else {
				fmt.Println("Error: cannot subtract non-number values")
				vm.push(types.NewNull())
			}

		case types.OpMultiply:
			right := vm.pop()
			left := vm.pop()

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				result := left.Data.(float64) * right.Data.(float64)
				vm.push(types.NewNumber(result))
			} else {
				fmt.Println("Error: cannot multiply non-number values")
				vm.push(types.NewNull())
			}

		case types.OpDivide:
			right := vm.pop()
			left := vm.pop()

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				// Check for division by zero
				if right.Data.(float64) == 0 {
					fmt.Println("Error: division by zero")
					vm.push(types.NewNull())
				} else {
					result := left.Data.(float64) / right.Data.(float64)
					vm.push(types.NewNumber(result))
				}
			} else {
				fmt.Println("Error: cannot divide non-number values")
				vm.push(types.NewNull())
			}

		case types.OpNegate:
			operand := vm.pop()

			if operand.Type == types.TypeNumber {
				result := -operand.Data.(float64)
				vm.push(types.NewNumber(result))
			} else {
				fmt.Println("Error: cannot negate non-number value")
				vm.push(types.NewNull())
			}

		// Comparison operators with safer implementation
		case types.OpEqual:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for equality comparison")
				vm.push(types.NewBoolean(false))
				continue
			}

			right := vm.pop()
			left := vm.pop()

			if left.Type != right.Type {
				vm.push(types.NewBoolean(false))
			} else {
				switch left.Type {
				case types.TypeNumber:
					vm.push(types.NewBoolean(left.Data.(float64) == right.Data.(float64)))
				case types.TypeString:
					vm.push(types.NewBoolean(left.Data.(string) == right.Data.(string)))
				case types.TypeBoolean:
					vm.push(types.NewBoolean(left.Data.(bool) == right.Data.(bool)))
				case types.TypeNull:
					vm.push(types.NewBoolean(true)) // null == null
				default:
					vm.push(types.NewBoolean(false))
				}
			}

		case types.OpNotEqual:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for inequality comparison")
				vm.push(types.NewBoolean(true))
				continue
			}

			right := vm.pop()
			left := vm.pop()

			if left.Type != right.Type {
				vm.push(types.NewBoolean(true))
			} else {
				switch left.Type {
				case types.TypeNumber:
					vm.push(types.NewBoolean(left.Data.(float64) != right.Data.(float64)))
				case types.TypeString:
					vm.push(types.NewBoolean(left.Data.(string) != right.Data.(string)))
				case types.TypeBoolean:
					vm.push(types.NewBoolean(left.Data.(bool) != right.Data.(bool)))
				case types.TypeNull:
					vm.push(types.NewBoolean(false)) // null != null is false
				default:
					vm.push(types.NewBoolean(true))
				}
			}

		case types.OpLessThan:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for less-than comparison")
				vm.push(types.NewBoolean(false))
				continue
			}

			// Peek at values first before popping
			right := vm.stack[vm.sp-1]
			left := vm.stack[vm.sp-2]

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				// Now pop them
				vm.pop()
				vm.pop()
				vm.push(types.NewBoolean(left.Data.(float64) < right.Data.(float64)))
			} else {
				// Pop the values to maintain stack balance
				vm.pop()
				vm.pop()
				fmt.Println("Error: cannot compare non-number values with <")
				vm.push(types.NewBoolean(false))
			}

		case types.OpGreaterThan:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for greater-than comparison")
				vm.push(types.NewBoolean(false))
				continue
			}

			// Peek at values first before popping
			right := vm.stack[vm.sp-1]
			left := vm.stack[vm.sp-2]

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				// Now pop them
				vm.pop()
				vm.pop()
				vm.push(types.NewBoolean(left.Data.(float64) > right.Data.(float64)))
			} else {
				// Pop the values to maintain stack balance
				vm.pop()
				vm.pop()
				fmt.Println("Error: cannot compare non-number values with >")
				vm.push(types.NewBoolean(false))
			}

		case types.OpLessEqual:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for less-equal comparison")
				vm.push(types.NewBoolean(false))
				continue
			}

			// Peek at values first before popping
			right := vm.stack[vm.sp-1]
			left := vm.stack[vm.sp-2]

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				// Now pop them
				vm.pop()
				vm.pop()
				vm.push(types.NewBoolean(left.Data.(float64) <= right.Data.(float64)))
			} else {
				// Pop the values to maintain stack balance
				vm.pop()
				vm.pop()
				fmt.Println("Error: cannot compare non-number values with <=")
				vm.push(types.NewBoolean(false))
			}

		case types.OpGreaterEqual:
			if vm.sp < 2 {
				fmt.Println("Error: not enough operands for greater-equal comparison")
				vm.push(types.NewBoolean(false))
				continue
			}

			// Peek at values first before popping
			right := vm.stack[vm.sp-1]
			left := vm.stack[vm.sp-2]

			if left.Type == types.TypeNumber && right.Type == types.TypeNumber {
				// Now pop them
				vm.pop()
				vm.pop()
				vm.push(types.NewBoolean(left.Data.(float64) >= right.Data.(float64)))
			} else {
				// Pop the values to maintain stack balance
				vm.pop()
				vm.pop()
				fmt.Println("Error: cannot compare non-number values with >=")
				vm.push(types.NewBoolean(false))
			}
		}
	}
}

func (vm *VM) push(value types.Value) {
	if vm.sp >= len(vm.stack) {
		// Grow stack if needed
		newStack := make([]types.Value, len(vm.stack)*2)
		copy(newStack, vm.stack)
		vm.stack = newStack
	}
	vm.stack[vm.sp] = value
	vm.sp++
}

func (vm *VM) pop() types.Value {
	if vm.sp <= 0 {
		// Return null instead of causing a panic when trying to pop from an empty stack
		fmt.Println("Stack underflow error")
		return types.NewNull()
	}
	vm.sp--
	return vm.stack[vm.sp]
}

func (vm *VM) formatTable(table *types.Table) string {
	result := "{"
	for i, entry := range table.Entries {
		result += vm.formatValue(entry.Key) + " = " + vm.formatValue(entry.Value)
		if i < len(table.Entries)-1 {
			result += ", "
		}
	}
	result += "}"
	return result
}

func (vm *VM) formatValue(value types.Value) string {
	switch value.Type {
	case types.TypeString:
		return "\"" + value.Data.(string) + "\""
	case types.TypeNumber:
		return fmt.Sprintf("%v", value.Data.(float64))
	case types.TypeBoolean:
		return fmt.Sprintf("%v", value.Data.(bool))
	case types.TypeNull:
		return "null"
	case types.TypeTable:
		return vm.formatTable(value.Data.(*types.Table))
	default:
		return "unknown"
	}
}