package compiler

import (
	"git.sharkk.net/Sharkk/Mako/parser"
	"git.sharkk.net/Sharkk/Mako/types"
)

// Compile converts AST to bytecode
func Compile(program *parser.Program) *types.Bytecode {
	c := &compiler{
		constants:    []any{},
		instructions: []types.Instruction{},
		scopes:       []scope{},
	}

	// Start in global scope
	c.enterScope()

	for _, stmt := range program.Statements {
		c.compileStatement(stmt)
	}

	c.exitScope()

	return &types.Bytecode{
		Constants:    c.constants,
		Instructions: c.instructions,
	}
}

type scope struct {
	variables map[string]bool
}

type compiler struct {
	constants    []any
	instructions []types.Instruction
	scopes       []scope
}

func (c *compiler) enterScope() {
	c.scopes = append(c.scopes, scope{
		variables: make(map[string]bool),
	})
	c.emit(types.OpEnterScope, 0)
}

func (c *compiler) exitScope() {
	c.scopes = c.scopes[:len(c.scopes)-1]
	c.emit(types.OpExitScope, 0)
}

func (c *compiler) declareVariable(name string) {
	if len(c.scopes) > 0 {
		c.scopes[len(c.scopes)-1].variables[name] = true
	}
}

func (c *compiler) isLocalVariable(name string) bool {
	for i := len(c.scopes) - 1; i >= 0; i-- {
		if _, ok := c.scopes[i].variables[name]; ok {
			return true
		}
	}
	return false
}

func (c *compiler) compileStatement(stmt parser.Statement) {
	switch s := stmt.(type) {
	case *parser.VariableStatement:
		c.compileExpression(s.Value)
		nameIndex := c.addConstant(s.Name.Value)

		// Use SetGlobal for top-level variables to persist between REPL lines
		if len(c.scopes) <= 1 {
			c.emit(types.OpSetGlobal, nameIndex)
		} else {
			c.declareVariable(s.Name.Value)
			c.emit(types.OpSetLocal, nameIndex)
		}

	case *parser.IndexAssignmentStatement:
		c.compileExpression(s.Left)
		c.compileExpression(s.Index)
		c.compileExpression(s.Value)
		c.emit(types.OpSetIndex, 0)

	case *parser.EchoStatement:
		c.compileExpression(s.Value)
		c.emit(types.OpEcho, 0)

	case *parser.BlockStatement:
		c.enterScope()
		for _, blockStmt := range s.Statements {
			c.compileStatement(blockStmt)
		}
		c.exitScope()
	}
}

func (c *compiler) compileExpression(expr parser.Expression) {
	switch e := expr.(type) {
	case *parser.StringLiteral:
		constIndex := c.addConstant(e.Value)
		c.emit(types.OpConstant, constIndex)

	case *parser.NumberLiteral:
		constIndex := c.addConstant(e.Value)
		c.emit(types.OpConstant, constIndex)

	case *parser.Identifier:
		nameIndex := c.addConstant(e.Value)

		// Check if it's a local variable first
		if c.isLocalVariable(e.Value) {
			c.emit(types.OpGetLocal, nameIndex)
		} else {
			// Otherwise treat as global
			c.emit(types.OpGetGlobal, nameIndex)
		}

	case *parser.TableLiteral:
		c.emit(types.OpNewTable, 0)

		for key, value := range e.Pairs {
			c.emit(types.OpDup, 0)

			// Special handling for identifier keys in tables
			if ident, ok := key.(*parser.Identifier); ok {
				// Treat identifiers as string literals in table keys
				strIndex := c.addConstant(ident.Value)
				c.emit(types.OpConstant, strIndex)
			} else {
				// For other expressions, compile normally
				c.compileExpression(key)
			}

			c.compileExpression(value)
			c.emit(types.OpSetIndex, 0)
			c.emit(types.OpPop, 0)
		}

	case *parser.IndexExpression:
		c.compileExpression(e.Left)
		c.compileExpression(e.Index)
		c.emit(types.OpGetIndex, 0)
	}
}

func (c *compiler) addConstant(value any) int {
	c.constants = append(c.constants, value)
	return len(c.constants) - 1
}

func (c *compiler) emit(op types.Opcode, operand int) {
	instruction := types.Instruction{
		Opcode:  op,
		Operand: operand,
	}
	c.instructions = append(c.instructions, instruction)
}