Mako/compiler/compiler.go

package compiler

import (
	"fmt"

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

// Compiler holds the compilation state and compiles AST to bytecode
type Compiler struct {
	current   *CompilerState // Current compilation state
	enclosing *CompilerState // Enclosing function state for closures
	errors    []CompileError // Compilation errors
}

// NewCompiler creates a new compiler instance
func NewCompiler() *Compiler {
	return &Compiler{
		current: NewCompilerState(FunctionTypeScript),
		errors:  make([]CompileError, 0),
	}
}

// Compile compiles a program AST to bytecode
func (c *Compiler) Compile(program *parser.Program) (*Chunk, []CompileError) {
	for _, stmt := range program.Statements {
		c.compileStatement(stmt)
	}

	c.current.EmitInstruction(OpReturnNil)

	if len(c.errors) > 0 {
		return nil, c.errors
	}

	return c.current.Chunk, nil
}

// Statement compilation
func (c *Compiler) compileStatement(stmt parser.Statement) {
	switch s := stmt.(type) {
	case *parser.StructStatement:
		c.compileStructStatement(s)
	case *parser.MethodDefinition:
		c.compileMethodDefinition(s)
	case *parser.Assignment:
		c.compileAssignment(s)
	case *parser.ExpressionStatement:
		c.compileExpression(s.Expression)
		c.current.EmitInstruction(OpPop) // Discard result
	case *parser.EchoStatement:
		c.compileExpression(s.Value)
		c.current.EmitInstruction(OpEcho)
	case *parser.IfStatement:
		c.compileIfStatement(s)
	case *parser.WhileStatement:
		c.compileWhileStatement(s)
	case *parser.ForStatement:
		c.compileForStatement(s)
	case *parser.ForInStatement:
		c.compileForInStatement(s)
	case *parser.ReturnStatement:
		c.compileReturnStatement(s)
	case *parser.ExitStatement:
		c.compileExitStatement(s)
	case *parser.BreakStatement:
		c.current.EmitBreak()
	default:
		c.addError(fmt.Sprintf("unknown statement type: %T", stmt))
	}
}

// Expression compilation
func (c *Compiler) compileExpression(expr parser.Expression) {
	switch e := expr.(type) {
	case *parser.Identifier:
		c.compileIdentifier(e)
	case *parser.NumberLiteral:
		c.compileNumberLiteral(e)
	case *parser.StringLiteral:
		c.compileStringLiteral(e)
	case *parser.BooleanLiteral:
		c.compileBooleanLiteral(e)
	case *parser.NilLiteral:
		c.compileNilLiteral(e)
	case *parser.TableLiteral:
		c.compileTableLiteral(e)
	case *parser.StructConstructor:
		c.compileStructConstructor(e)
	case *parser.FunctionLiteral:
		c.compileFunctionLiteral(e)
	case *parser.CallExpression:
		c.compileCallExpression(e)
	case *parser.PrefixExpression:
		c.compilePrefixExpression(e)
	case *parser.InfixExpression:
		c.compileInfixExpression(e)
	case *parser.IndexExpression:
		c.compileIndexExpression(e)
	case *parser.DotExpression:
		c.compileDotExpression(e)
	case *parser.Assignment:
		c.compileAssignmentExpression(e)
	default:
		c.addError(fmt.Sprintf("unknown expression type: %T", expr))
	}
}

// Literal compilation
func (c *Compiler) compileNumberLiteral(node *parser.NumberLiteral) {
	value := Value{Type: ValueNumber, Data: node.Value}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpLoadConst, uint16(index))
}

func (c *Compiler) compileStringLiteral(node *parser.StringLiteral) {
	value := Value{Type: ValueString, Data: node.Value}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpLoadConst, uint16(index))
}

func (c *Compiler) compileBooleanLiteral(node *parser.BooleanLiteral) {
	value := Value{Type: ValueBool, Data: node.Value}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpLoadConst, uint16(index))
}

func (c *Compiler) compileNilLiteral(node *parser.NilLiteral) {
	value := Value{Type: ValueNil, Data: nil}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpLoadConst, uint16(index))
}

// Identifier compilation
func (c *Compiler) compileIdentifier(node *parser.Identifier) {
	// Try local variables first
	slot := c.current.ResolveLocal(node.Value)
	if slot != -1 {
		if slot == -2 {
			c.addError("can't read local variable in its own initializer")
			return
		}
		c.current.EmitInstruction(OpLoadLocal, uint16(slot))
		return
	}

	// Try upvalues
	upvalue := c.resolveUpvalue(node.Value)
	if upvalue != -1 {
		c.current.EmitInstruction(OpGetUpvalue, uint16(upvalue))
		return
	}

	// Must be global
	value := Value{Type: ValueString, Data: node.Value}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpLoadGlobal, uint16(index))
}

// Assignment compilation
func (c *Compiler) compileAssignment(node *parser.Assignment) {
	c.compileExpression(node.Value)

	switch target := node.Target.(type) {
	case *parser.Identifier:
		if node.IsDeclaration {
			// Check if we're at global scope
			if c.current.FunctionType == FunctionTypeScript && c.current.ScopeDepth == 0 {
				// Global variable declaration - treat as global assignment
				value := Value{Type: ValueString, Data: target.Value}
				index := c.current.AddConstant(value)
				if index == -1 {
					c.addError("too many constants")
					return
				}
				c.current.EmitInstruction(OpStoreGlobal, uint16(index))
			} else {
				// Local variable declaration
				if err := c.current.AddLocal(target.Value); err != nil {
					c.addError(err.Error())
					return
				}
				c.current.MarkInitialized()
			}
		} else {
			// Assignment to existing variable
			slot := c.current.ResolveLocal(target.Value)
			if slot != -1 {
				c.current.EmitInstruction(OpStoreLocal, uint16(slot))
			} else {
				upvalue := c.resolveUpvalue(target.Value)
				if upvalue != -1 {
					c.current.EmitInstruction(OpSetUpvalue, uint16(upvalue))
				} else {
					// Global assignment
					value := Value{Type: ValueString, Data: target.Value}
					index := c.current.AddConstant(value)
					if index == -1 {
						c.addError("too many constants")
						return
					}
					c.current.EmitInstruction(OpStoreGlobal, uint16(index))
				}
			}
		}
	case *parser.DotExpression:
		// table.field = value
		c.compileExpression(target.Left)
		value := Value{Type: ValueString, Data: target.Key}
		index := c.current.AddConstant(value)
		if index == -1 {
			c.addError("too many constants")
			return
		}
		c.current.EmitInstruction(OpSetField, uint16(index))
	case *parser.IndexExpression:
		// table[key] = value
		c.compileExpression(target.Left)
		c.compileExpression(target.Index)
		c.current.EmitInstruction(OpSetIndex)
	default:
		c.addError("invalid assignment target")
	}
}

func (c *Compiler) compileAssignmentExpression(node *parser.Assignment) {
	c.compileAssignment(node)
	// Assignment expressions leave the assigned value on stack
}

// Operator compilation
func (c *Compiler) compilePrefixExpression(node *parser.PrefixExpression) {
	c.compileExpression(node.Right)

	switch node.Operator {
	case "-":
		c.current.EmitInstruction(OpNeg)
	case "not":
		c.current.EmitInstruction(OpNot)
	default:
		c.addError(fmt.Sprintf("unknown prefix operator: %s", node.Operator))
	}
}

func (c *Compiler) compileInfixExpression(node *parser.InfixExpression) {
	// Handle short-circuit operators specially
	if node.Operator == "and" {
		c.compileExpression(node.Left)
		jump := c.current.EmitJump(OpJumpIfFalse)
		c.current.EmitInstruction(OpPop)
		c.compileExpression(node.Right)
		c.current.PatchJump(jump)
		return
	}

	if node.Operator == "or" {
		c.compileExpression(node.Left)
		elseJump := c.current.EmitJump(OpJumpIfFalse)
		endJump := c.current.EmitJump(OpJump)
		c.current.PatchJump(elseJump)
		c.current.EmitInstruction(OpPop)
		c.compileExpression(node.Right)
		c.current.PatchJump(endJump)
		return
	}

	// Regular binary operators
	c.compileExpression(node.Left)
	c.compileExpression(node.Right)

	switch node.Operator {
	case "+":
		c.current.EmitInstruction(OpAdd)
	case "-":
		c.current.EmitInstruction(OpSub)
	case "*":
		c.current.EmitInstruction(OpMul)
	case "/":
		c.current.EmitInstruction(OpDiv)
	case "==":
		c.current.EmitInstruction(OpEq)
	case "!=":
		c.current.EmitInstruction(OpNeq)
	case "<":
		c.current.EmitInstruction(OpLt)
	case "<=":
		c.current.EmitInstruction(OpLte)
	case ">":
		c.current.EmitInstruction(OpGt)
	case ">=":
		c.current.EmitInstruction(OpGte)
	default:
		c.addError(fmt.Sprintf("unknown infix operator: %s", node.Operator))
	}
}

// Control flow compilation
func (c *Compiler) compileIfStatement(node *parser.IfStatement) {
	c.compileExpression(node.Condition)

	// Jump over then branch if condition is false
	thenJump := c.current.EmitJump(OpJumpIfFalse)
	c.current.EmitInstruction(OpPop)

	// Compile then branch
	c.current.BeginScope()
	for _, stmt := range node.Body {
		c.compileStatement(stmt)
	}
	c.current.EndScope()

	// Jump over else branches
	elseJump := c.current.EmitJump(OpJump)
	c.current.PatchJump(thenJump)
	c.current.EmitInstruction(OpPop)

	// Compile elseif branches
	var elseifJumps []int
	for _, elseif := range node.ElseIfs {
		c.compileExpression(elseif.Condition)
		nextJump := c.current.EmitJump(OpJumpIfFalse)
		c.current.EmitInstruction(OpPop)

		c.current.BeginScope()
		for _, stmt := range elseif.Body {
			c.compileStatement(stmt)
		}
		c.current.EndScope()

		elseifJumps = append(elseifJumps, c.current.EmitJump(OpJump))
		c.current.PatchJump(nextJump)
		c.current.EmitInstruction(OpPop)
	}

	// Compile else branch
	if len(node.Else) > 0 {
		c.current.BeginScope()
		for _, stmt := range node.Else {
			c.compileStatement(stmt)
		}
		c.current.EndScope()
	}

	// Patch all jumps to end
	c.current.PatchJump(elseJump)
	for _, jump := range elseifJumps {
		c.current.PatchJump(jump)
	}
}

func (c *Compiler) compileWhileStatement(node *parser.WhileStatement) {
	c.current.EnterLoop()

	c.compileExpression(node.Condition)
	exitJump := c.current.EmitJump(OpJumpIfFalse)
	c.current.EmitInstruction(OpPop)

	c.current.BeginScope()
	for _, stmt := range node.Body {
		c.compileStatement(stmt)
	}
	c.current.EndScope()

	// Jump back to condition
	jump := len(c.current.Chunk.Code) - c.current.LoopStart + 2
	c.current.EmitInstruction(OpJump, uint16(jump))

	c.current.PatchJump(exitJump)
	c.current.EmitInstruction(OpPop)

	c.current.ExitLoop()
}

// Table operations
func (c *Compiler) compileTableLiteral(node *parser.TableLiteral) {
	c.current.EmitInstruction(OpNewTable)

	for _, pair := range node.Pairs {
		if pair.Key == nil {
			// Array-style element
			c.compileExpression(pair.Value)
			c.current.EmitInstruction(OpTableInsert)
		} else {
			// Key-value pair
			c.current.EmitInstruction(OpDup) // Duplicate table reference
			c.compileExpression(pair.Key)
			c.compileExpression(pair.Value)
			c.current.EmitInstruction(OpSetIndex)
		}
	}
}

func (c *Compiler) compileDotExpression(node *parser.DotExpression) {
	c.compileExpression(node.Left)
	value := Value{Type: ValueString, Data: node.Key}
	index := c.current.AddConstant(value)
	if index == -1 {
		c.addError("too many constants")
		return
	}
	c.current.EmitInstruction(OpGetField, uint16(index))
}

func (c *Compiler) compileIndexExpression(node *parser.IndexExpression) {
	c.compileExpression(node.Left)
	c.compileExpression(node.Index)
	c.current.EmitInstruction(OpGetIndex)
}

// Function compilation
func (c *Compiler) compileCallExpression(node *parser.CallExpression) {
	c.compileExpression(node.Function)

	// Compile arguments
	for _, arg := range node.Arguments {
		c.compileExpression(arg)
	}

	c.current.EmitInstruction(OpCall, uint16(len(node.Arguments)))
}

func (c *Compiler) compileReturnStatement(node *parser.ReturnStatement) {
	if node.Value != nil {
		c.compileExpression(node.Value)
		c.current.EmitInstruction(OpReturn)
	} else {
		c.current.EmitInstruction(OpReturnNil)
	}
}

func (c *Compiler) compileExitStatement(node *parser.ExitStatement) {
	if node.Value != nil {
		c.compileExpression(node.Value)
	} else {
		// Default exit code 0
		value := Value{Type: ValueNumber, Data: float64(0)}
		index := c.current.AddConstant(value)
		if index == -1 {
			c.addError("too many constants")
			return
		}
		c.current.EmitInstruction(OpLoadConst, uint16(index))
	}
	c.current.EmitInstruction(OpExit)
}

// Placeholder implementations for complex features
func (c *Compiler) compileStructStatement(node *parser.StructStatement) {
	// TODO: Implement struct compilation
	c.addError("struct compilation not yet implemented")
}

func (c *Compiler) compileMethodDefinition(node *parser.MethodDefinition) {
	// TODO: Implement method compilation
	c.addError("method compilation not yet implemented")
}

func (c *Compiler) compileStructConstructor(node *parser.StructConstructor) {
	// TODO: Implement struct constructor compilation
	c.addError("struct constructor compilation not yet implemented")
}

func (c *Compiler) compileFunctionLiteral(node *parser.FunctionLiteral) {
	// TODO: Implement function literal compilation
	c.addError("function literal compilation not yet implemented")
}

func (c *Compiler) compileForStatement(node *parser.ForStatement) {
	// TODO: Implement numeric for loop compilation
	c.addError("for statement compilation not yet implemented")
}

func (c *Compiler) compileForInStatement(node *parser.ForInStatement) {
	// TODO: Implement for-in loop compilation
	c.addError("for-in statement compilation not yet implemented")
}

// Helper methods
func (c *Compiler) resolveUpvalue(name string) int {
	if c.enclosing == nil {
		return -1
	}

	local := c.enclosing.ResolveLocal(name)
	if local != -1 {
		c.enclosing.Locals[local].IsCaptured = true
		return c.current.AddUpvalue(uint8(local), true)
	}

	upvalue := c.resolveUpvalueInEnclosing(name)
	if upvalue != -1 {
		return c.current.AddUpvalue(uint8(upvalue), false)
	}

	return -1
}

func (c *Compiler) resolveUpvalueInEnclosing(name string) int {
	if c.enclosing == nil {
		return -1
	}

	// This would recursively check enclosing scopes
	// Simplified for now
	return -1
}

func (c *Compiler) addError(message string) {
	c.errors = append(c.errors, CompileError{
		Message: message,
		Line:    0, // TODO: Add line tracking
		Column:  0, // TODO: Add column tracking
	})
}

// Error reporting
func (c *Compiler) Errors() []CompileError { return c.errors }
func (c *Compiler) HasErrors() bool        { return len(c.errors) > 0 }