c65gm/internal/commands/for.go

package commands

import (
	"fmt"
	"os"
	"strings"

	"c65gm/internal/compiler"
	"c65gm/internal/preproc"
	"c65gm/internal/utils"
)

// ForCommand handles FOR loop statements
// Syntax: FOR <var> = <start> TO <end> [STEP <step>]
type ForCommand struct {
	varName     string
	varKind     compiler.VarKind
	startOp     *compiler.OperandInfo
	endOp       *compiler.OperandInfo
	stepOp      *compiler.OperandInfo
	useLongJump bool
	loopLabel   string
	skipLabel   string
}

func (c *ForCommand) WillHandle(line preproc.Line) bool {
	params, err := utils.ParseParams(line.Text)
	if err != nil || len(params) == 0 {
		return false
	}
	return strings.ToUpper(params[0]) == "FOR"
}

func (c *ForCommand) Interpret(line preproc.Line, ctx *compiler.CompilerContext) error {
	params, err := utils.ParseParams(line.Text)
	if err != nil {
		return err
	}

	// FOR <var> = <start> TO/DOWNTO <end> [STEP <step>]
	// Minimum: 6 params (FOR var = start TO end)
	// Maximum: 8 params (FOR var = start TO end STEP step)
	if len(params) < 6 { // FOR keyword goes towards count
		return fmt.Errorf("FOR: expected at least 5 parameters, got %d", len(params))
	}

	if len(params) != 6 && len(params) != 8 {
		return fmt.Errorf("FOR: expected 5 or 7 parameters, got %d", len(params))
	}

	// Check '=' separator
	if params[2] != "=" {
		return fmt.Errorf("FOR: expected '=' at position 3, got %q", params[2])
	}

	scope := ctx.CurrentScope()
	constLookup := func(name string) (int64, bool) {
		sym := ctx.SymbolTable.Lookup(name, scope)
		if sym != nil && sym.IsConst() {
			return int64(sym.Value), true
		}
		return 0, false
	}

	// Parse variable
	varName := params[1]
	varSym := ctx.SymbolTable.Lookup(varName, scope)
	if varSym == nil {
		return fmt.Errorf("FOR: unknown variable %q", varName)
	}
	if varSym.IsConst() {
		return fmt.Errorf("FOR: cannot use constant %q as loop variable", varName)
	}
	c.varName = varSym.FullName()
	c.varKind = varSym.GetVarKind()

	// Parse start value
	var parseErr error
	startVarName, startVarKind, startValue, startIsVar, parseErr := compiler.ParseOperandParam(
		params[3], ctx.SymbolTable, scope, constLookup)
	if parseErr != nil {
		return fmt.Errorf("FOR: start value: %w", parseErr)
	}
	c.startOp = &compiler.OperandInfo{
		VarName: startVarName,
		VarKind: startVarKind,
		Value:   startValue,
		IsVar:   startIsVar,
	}

	// Parse direction (TO only)
	direction := strings.ToUpper(params[4])
	if direction != "TO" {
		return fmt.Errorf("FOR: expected 'TO' at position 5, got %q (DOWNTO is not supported)", params[4])
	}

	// Parse end value
	endVarName, endVarKind, endValue, endIsVar, parseErr := compiler.ParseOperandParam(
		params[5], ctx.SymbolTable, scope, constLookup)
	if parseErr != nil {
		return fmt.Errorf("FOR: end value: %w", parseErr)
	}
	c.endOp = &compiler.OperandInfo{
		VarName: endVarName,
		VarKind: endVarKind,
		Value:   endValue,
		IsVar:   endIsVar,
	}

	if c.varKind == compiler.KindByte {
		// Error on literal out of range
		if !c.startOp.IsVar && c.startOp.Value > 255 {
			return fmt.Errorf("FOR: BYTE variable cannot start at literal %d (max 255)", c.startOp.Value)
		}
		if !c.endOp.IsVar && c.endOp.Value > 255 {
			return fmt.Errorf("FOR: BYTE variable cannot loop to literal %d (max 255)", c.endOp.Value)
		}

		// Warn on variable type mismatch
		if c.startOp.IsVar && c.startOp.VarKind == compiler.KindWord {
			_, _ = fmt.Fprintf(os.Stderr, "%s:%d: warning: BYTE loop variable with WORD start value truncates to low byte\n",
				line.Filename, line.LineNo)
		}
		if c.endOp.IsVar && c.endOp.VarKind == compiler.KindWord {
			_, _ = fmt.Fprintf(os.Stderr, "%s:%d: warning: BYTE loop variable with WORD end value may cause infinite loop\n",
				line.Filename, line.LineNo)
		}
	}

	// Parse optional STEP
	if len(params) == 8 {
		if strings.ToUpper(params[6]) != "STEP" {
			return fmt.Errorf("FOR: expected 'STEP' at position 7, got %q", params[6])
		}

		stepVarName, stepVarKind, stepValue, stepIsVar, parseErr := compiler.ParseOperandParam(
			params[7], ctx.SymbolTable, scope, constLookup)
		if parseErr != nil {
			return fmt.Errorf("FOR: step value: %w", parseErr)
		}

		// Check for zero or negative step if literal
		if !stepIsVar {
			if stepValue == 0 {
				return fmt.Errorf("FOR: STEP cannot be zero")
			}
			// Since BYTE and WORD are unsigned, values > 32767 are treated as large positive
			// We don't allow negative literals since they'd be interpreted as large unsigned
			// This is a reasonable restriction for step values
		}

		c.stepOp = &compiler.OperandInfo{
			VarName: stepVarName,
			VarKind: stepVarKind,
			Value:   stepValue,
			IsVar:   stepIsVar,
		}
	} else {
		// Default STEP 1
		c.stepOp = &compiler.OperandInfo{
			Value: 1,
			IsVar: false,
		}
	}

	// Check pragma
	ps := ctx.Pragma.GetPragmaSetByIndex(line.PragmaSetIndex)
	longJumpPragma := ps.GetPragma("_P_USE_LONG_JUMP")
	c.useLongJump = longJumpPragma != "" && longJumpPragma != "0"

	// Create labels
	c.loopLabel = ctx.LoopStartStack.Push()
	c.skipLabel = ctx.LoopEndStack.Push()

	// Push FOR info to ForStack
	ctx.ForStack.Push(&compiler.ForLoopInfo{
		VarName:     c.varName,
		VarKind:     c.varKind,
		EndOperand:  c.endOp,
		StepOperand: c.stepOp,
		LoopLabel:   c.loopLabel,
		SkipLabel:   c.skipLabel,
	})

	return nil
}

func (c *ForCommand) Generate(ctx *compiler.CompilerContext) ([]string, error) {
	var asm []string

	// Initial assignment: var = start
	assignAsm := c.generateAssignment()
	asm = append(asm, assignAsm...)

	// Emit loop label
	asm = append(asm, c.loopLabel)

	// Generate comparison for TO loop: continue if var <= end (skip if var > end)
	varOp := &operandInfo{
		varName: c.varName,
		varKind: c.varKind,
		isVar:   true,
	}

	// Convert compiler.OperandInfo to commands.operandInfo for comparison
	endOp := &operandInfo{
		varName: c.endOp.VarName,
		varKind: c.endOp.VarKind,
		value:   c.endOp.Value,
		isVar:   c.endOp.IsVar,
	}

	gen, err := newComparisonGenerator(
		opLessEqual,
		varOp,
		endOp,
		c.useLongJump,
		ctx.LoopEndStack,
		ctx.GeneralStack,
	)
	if err != nil {
		return nil, fmt.Errorf("FOR: %w", err)
	}

	cmpAsm, err := gen.generate()
	if err != nil {
		return nil, fmt.Errorf("FOR: %w", err)
	}

	asm = append(asm, cmpAsm...)
	return asm, nil
}

func (c *ForCommand) generateAssignment() []string {
	var asm []string

	// Variable assignment from startOp
	if c.startOp.IsVar {
		// Destination: byte
		if c.varKind == compiler.KindByte {
			// byte → byte or word → byte (take low byte)
			asm = append(asm, fmt.Sprintf("\tlda %s", c.startOp.VarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.varName))
			return asm
		}

		// Destination: word
		// byte → word (zero-extend)
		if c.startOp.VarKind == compiler.KindByte {
			asm = append(asm, fmt.Sprintf("\tlda %s", c.startOp.VarName))
			asm = append(asm, fmt.Sprintf("\tsta %s", c.varName))
			asm = append(asm, "\tlda #0")
			asm = append(asm, fmt.Sprintf("\tsta %s+1", c.varName))
			return asm
		}

		// word → word (copy both bytes)
		asm = append(asm, fmt.Sprintf("\tlda %s", c.startOp.VarName))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.varName))
		asm = append(asm, fmt.Sprintf("\tlda %s+1", c.startOp.VarName))
		asm = append(asm, fmt.Sprintf("\tsta %s+1", c.varName))
		return asm
	}

	// Literal assignment
	lo := uint8(c.startOp.Value & 0xFF)
	hi := uint8((c.startOp.Value >> 8) & 0xFF)

	// Destination: byte
	if c.varKind == compiler.KindByte {
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", lo))
		asm = append(asm, fmt.Sprintf("\tsta %s", c.varName))
		return asm
	}

	// Destination: word
	asm = append(asm, fmt.Sprintf("\tlda #$%02x", lo))
	asm = append(asm, fmt.Sprintf("\tsta %s", c.varName))

	// Optimization: don't reload if lo == hi
	if lo != hi {
		asm = append(asm, fmt.Sprintf("\tlda #$%02x", hi))
	}
	asm = append(asm, fmt.Sprintf("\tsta %s+1", c.varName))

	return asm
}