package vm

import (
	"fmt"

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

type Opcode byte

const (
	OpConstant  Opcode = iota
	OpSetLocal         // Set local variable
	OpGetLocal         // Get local variable
	OpSetGlobal        // Set global variable
	OpGetGlobal        // Get global variable
	OpEcho
	OpNewTable
	OpSetIndex
	OpGetIndex
	OpDup
	OpPop
	OpEnterScope // New opcode for entering a block
	OpExitScope  // New opcode for exiting a block
)

type Instruction struct {
	Opcode  Opcode
	Operand int
}

type Bytecode struct {
	Constants    []any
	Instructions []Instruction
}

// 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 *Bytecode) {
	vm.constants = bytecode.Constants

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

		switch instruction.Opcode {
		case 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 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 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 OpSetGlobal:
			constIndex := instruction.Operand
			name := vm.constants[constIndex].(string)
			value := vm.pop()
			vm.globals[name] = value

		case 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 OpEnterScope:
			// Push a new scope
			vm.scopes = append(vm.scopes, Scope{
				Variables: make(map[string]types.Value),
			})

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

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

		case 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 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 OpDup:
			if vm.sp > 0 {
				vm.push(vm.stack[vm.sp-1])
			}

		case OpPop:
			vm.pop()

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

func (vm *VM) push(value types.Value) {
	vm.stack[vm.sp] = value
	vm.sp++
}

func (vm *VM) pop() types.Value {
	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"
	}
}