c65gm/internal/compiler/scriptexec.go

package compiler

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"strings"

	"c65gm/internal/preproc"
	"c65gm/internal/utils"

	"go.starlark.net/lib/math"
	"go.starlark.net/starlark"
)

// mapStarlarkLine maps a Starlark 1-based line number to an index into scriptLines
// (a slice of preproc.Line). Returns -1 if the line cannot be mapped.
// For non-library scripts, the Starlark source is:
//
//	Line 1: def _main():
//	Line 2..N+1: indented script lines
//	Line N+2: _main()
//
// For library scripts, lines are used as-is.
func mapStarlarkLine(starlarkLine int, numScriptLines int, isLibrary bool) int {
	var idx int
	if isLibrary {
		idx = starlarkLine - 1
	} else {
		idx = starlarkLine - 2
	}
	if idx < 0 || idx >= numScriptLines {
		return -1
	}
	return idx
}

// starlarkPosition extracts the source position from a Starlark error.
// Returns the 1-based line number (0 if unknown).
func starlarkPosition(err error) int {
	var evalErr *starlark.EvalError
	if errors.As(err, &evalErr) {
		if len(evalErr.CallStack) > 0 {
			top := evalErr.CallStack.At(0)
			return int(top.Pos.Line)
		}
	}
	return 0
}

// starlarkErrorMsg extracts just the message from a Starlark error.
func starlarkErrorMsg(err error) string {
	var evalErr *starlark.EvalError
	if errors.As(err, &evalErr) {
		return evalErr.Msg
	}
	return err.Error()
}

// printScriptErrorContext prints a Starlark error with script-block-bounded source context,
// synthesizing SCRIPT/ENDSCRIPT boundary lines that the preprocessor discards.
// When the original Starlark error has a call stack with frames beyond the script block
// (e.g. a library function), those frames are shown as a backtrace below the source context.
func printScriptErrorContext(err error, errMsg string, scriptLines []preproc.Line, errorIdx int, blockType string) {
	if len(scriptLines) == 0 || errorIdx < 0 || errorIdx >= len(scriptLines) {
		fmt.Fprintf(os.Stderr, "\nError: Starlark error: %s\n\n", errMsg)
		return
	}

	line := scriptLines[errorIdx]
	filename := line.Filename

	const contextLines = 3

	fmt.Fprintf(os.Stderr, "\nError: Starlark error: %s\n", errMsg)
	fmt.Fprintf(os.Stderr, "  --> %s:%d\n\n", filename, line.LineNo)

	startIdx := errorIdx - contextLines
	if startIdx < 0 {
		startIdx = 0
	}
	endIdx := errorIdx + contextLines
	if endIdx >= len(scriptLines) {
		endIdx = len(scriptLines) - 1
	}

	scriptMarkerLineNo := scriptLines[0].LineNo - 1
	endScriptLineNo := scriptLines[len(scriptLines)-1].LineNo + 1

	maxLineNo := endScriptLineNo
	if scriptLines[endIdx].LineNo > maxLineNo {
		maxLineNo = scriptLines[endIdx].LineNo
	}
	if scriptMarkerLineNo > maxLineNo {
		maxLineNo = scriptMarkerLineNo
	}
	lineNumWidth := len(fmt.Sprintf("%d", maxLineNo))

	if startIdx == 0 {
		fmt.Fprintf(os.Stderr, "   %*d | %s\n", lineNumWidth, scriptMarkerLineNo, blockType)
	}

	for i := startIdx; i <= endIdx; i++ {
		l := scriptLines[i]
		marker := "   "
		if i == errorIdx {
			marker = ">> "
		}
		fmt.Fprintf(os.Stderr, "%s%*d | %s\n", marker, lineNumWidth, l.LineNo, l.Text)
	}

	if endIdx == len(scriptLines)-1 {
		fmt.Fprintf(os.Stderr, "   %*d | END%s\n", lineNumWidth, endScriptLineNo, blockType)
	}

	// Show call stack if the error has additional frames beyond the script block
	printStarlarkBacktrace(err, line.LineNo, filename)

	fmt.Fprintf(os.Stderr, "\n")
}

// printStarlarkBacktrace prints a concise call stack from a Starlark error,
// filtering out frames from the given scriptBlockLine and <toplevel>/_main wrappers.
func printStarlarkBacktrace(err error, scriptBlockLine int, scriptFile string) {
	var evalErr *starlark.EvalError
	if !errors.As(err, &evalErr) || len(evalErr.CallStack) < 2 {
		return
	}

	var frames []string
	for _, cf := range evalErr.CallStack {
		name := cf.Name
		pos := cf.Pos
		// Skip frames that are already shown in the bounded source context
		if name == "<toplevel>" || name == "_main" {
			continue
		}
		if int(pos.Line) == scriptBlockLine && pos.Filename() == scriptFile {
			continue
		}
		if pos.Filename() != "" && int(pos.Line) > 0 {
			frames = append(frames, fmt.Sprintf("  %s: in %s", pos.String(), name))
		}
	}
	if len(frames) > 0 {
		fmt.Fprintf(os.Stderr, "Call stack:\n%s\n", strings.Join(frames, "\n"))
	}
}

// printScriptErrorFallback prints what we can from a Starlark error when
// we couldn't map it to a specific source line within a script block.
func printScriptErrorFallback(err error) {
	fmt.Fprintf(os.Stderr, "\nError: Starlark error: %s\n", err)
	var evalErr *starlark.EvalError
	if errors.As(err, &evalErr) && len(evalErr.CallStack) > 0 {
		fmt.Fprintf(os.Stderr, "Call stack:\n")
		for _, cf := range evalErr.CallStack {
			name := cf.Name
			pos := cf.Pos
			if pos.Filename() != "" && int(pos.Line) > 0 {
				fmt.Fprintf(os.Stderr, "  %s: in %s\n", pos.String(), name)
			}
		}
	}
	fmt.Fprintf(os.Stderr, "\n")
}

// executeScript runs a Starlark script and returns the output lines.
// If isLibrary is true, the script is executed at top level (no _main wrapper)
// and resulting globals are persisted to ctx.ScriptLibraryGlobals.
func executeScript(scriptLines []preproc.Line, ctx *CompilerContext, isLibrary bool) ([]string, error) {
	// Extract text from preproc.Lines
	texts := make([]string, len(scriptLines))
	for i, l := range scriptLines {
		texts[i] = l.Text
	}

	// Join script lines
	scriptText := strings.Join(texts, "\n")

	// Expand |varname| -> actual variable names
	scriptText = expandVariables(scriptText, ctx)

	// Determine the source filename for Starlark
	sourceFile := scriptLines[0].Filename

	var finalScript string
	var starlarkFilename string
	if isLibrary {
		// LIBRARY: execute at top level so defs become globals
		finalScript = scriptText
		starlarkFilename = sourceFile
	} else {
		// Regular SCRIPT: wrap in function (Starlark requires control flow inside functions)
		finalScript = "def _main():\n"
		for _, line := range strings.Split(scriptText, "\n") {
			finalScript += "    " + line + "\n"
		}
		finalScript += "_main()\n"
		starlarkFilename = sourceFile
	}

	// Capture print output
	var output bytes.Buffer
	thread := &starlark.Thread{
		Print: func(_ *starlark.Thread, msg string) {
			output.WriteString(msg)
			output.WriteString("\n")
		},
	}

	// Set execution limit (prevent infinite loops)
	thread.SetMaxExecutionSteps(1000000) // 1M steps

	// Build predeclared: math module + library globals
	predeclared := starlark.StringDict{
		"math": math.Module,
	}
	for k, v := range ctx.ScriptLibraryGlobals {
		predeclared[k] = v
	}

	// Execute
	globals, err := starlark.ExecFile(thread, starlarkFilename, finalScript, predeclared)
	if err != nil {
		// Map Starlark error position back to source
		starLine := starlarkPosition(err)
		if starLine > 0 {
			idx := mapStarlarkLine(starLine, len(scriptLines), isLibrary)
			if idx >= 0 {
				msg := starlarkErrorMsg(err)
				blockType := "SCRIPT"
				if isLibrary {
					blockType = "SCRIPT LIBRARY"
				}
				printScriptErrorContext(err, msg, scriptLines, idx, blockType)
				return nil, fmt.Errorf("Starlark error: %s", msg)
			}
		}
		// Fallback: print whatever info we can extract
		printScriptErrorFallback(err)
		return nil, fmt.Errorf("Starlark error: %w", err)
	}

	// For LIBRARY: persist new globals (functions, variables defined at top level)
	if isLibrary {
		for k, v := range globals {
			ctx.ScriptLibraryGlobals[k] = v
		}
	}

	// Split output into lines for assembly
	outputStr := output.String()
	if outputStr == "" {
		return []string{}, nil
	}
	lines := strings.Split(strings.TrimRight(outputStr, "\n"), "\n")
	return lines, nil
}

// expandVariables replaces |varname| with expanded variable names from symbol table
func expandVariables(text string, ctx *CompilerContext) string {
	result := text
	for {
		start := strings.IndexByte(result, '|')
		if start == -1 {
			break
		}
		end := strings.IndexByte(result[start+1:], '|')
		if end == -1 {
			break // unclosed, let script fail
		}
		end += start + 1

		varName := result[start+1 : end]
		expandedName := ctx.SymbolTable.ExpandName(varName, ctx.CurrentScope())
		result = result[:start] + expandedName + result[end+1:]
	}
	return result
}

// ExecuteMacro executes a named macro with the given arguments and returns output lines
func ExecuteMacro(macroName string, args []string, ctx *CompilerContext) ([]string, error) {
	// Look up the macro
	macro, ok := ctx.ScriptMacros[macroName]
	if !ok {
		return nil, fmt.Errorf("undefined macro: %s", macroName)
	}

	// Check argument count
	if len(args) != len(macro.Params) {
		return nil, fmt.Errorf("macro %s expects %d arguments, got %d", macroName, len(macro.Params), len(args))
	}

	// Evaluate arguments and build parameter bindings
	paramBindings := make(starlark.StringDict)
	for i, arg := range args {
		val, err := evaluateMacroArg(arg, ctx)
		if err != nil {
			return nil, fmt.Errorf("error evaluating argument %d for macro %s: %w", i+1, macroName, err)
		}
		paramBindings[macro.Params[i]] = val
	}

	// Build the script: wrap macro body in a function with parameters bound
	scriptText := strings.Join(macro.Body, "\n")

	// Wrap in function for control flow support
	finalScript := "def _macro():\n"
	for _, line := range strings.Split(scriptText, "\n") {
		finalScript += "    " + line + "\n"
	}
	finalScript += "_macro()\n"

	// Use the source file where the macro was defined
	starlarkFilename := macro.SourceFile
	if starlarkFilename == "" {
		starlarkFilename = "macro.star"
	}

	// Capture print output
	var output bytes.Buffer
	thread := &starlark.Thread{
		Print: func(_ *starlark.Thread, msg string) {
			output.WriteString(msg)
			output.WriteString("\n")
		},
	}

	// Set execution limit
	thread.SetMaxExecutionSteps(1000000)

	// Build predeclared: math + library globals + parameter bindings
	predeclared := starlark.StringDict{
		"math": math.Module,
	}
	for k, v := range ctx.ScriptLibraryGlobals {
		predeclared[k] = v
	}
	for k, v := range paramBindings {
		predeclared[k] = v
	}

	// Execute
	_, err := starlark.ExecFile(thread, starlarkFilename, finalScript, predeclared)
	if err != nil {
		// Map error position back to macro definition site
		starLine := starlarkPosition(err)
		if starLine > 0 && macro.SourceFile != "" {
			idx := mapStarlarkLine(starLine, len(macro.Body), false) // macro is always wrapped
			if idx >= 0 {
				sourceLine := macro.StartLine + idx
				msg := starlarkErrorMsg(err)
				return nil, fmt.Errorf("Starlark error: at %s:%d: %s", macro.SourceFile, sourceLine, msg)
			}
		}
		return nil, fmt.Errorf("Starlark error: %w", err)
	}

	// Split output into lines
	outputStr := output.String()
	if outputStr == "" {
		return []string{}, nil
	}

	// Expand |varname| -> actual variable names in the OUTPUT
	// This happens at call site, so local variables are resolved using caller's scope
	outputStr = expandVariables(outputStr, ctx)

	return strings.Split(strings.TrimRight(outputStr, "\n"), "\n"), nil
}

// evaluateMacroArg evaluates a macro argument, returning either an int or string Starlark value
func evaluateMacroArg(arg string, ctx *CompilerContext) (starlark.Value, error) {
	arg = strings.TrimSpace(arg)

	// Create lookup function for constants
	lookup := ctx.SymbolTable.ConstantLookupFunc(nil)

	// Try to evaluate as expression (number, constant, arithmetic)
	val, err := utils.EvaluateExpression(arg, lookup)
	if err == nil {
		// Successfully evaluated as integer
		return starlark.MakeInt64(val), nil
	}

	// If it's a valid identifier, treat as label (string)
	if utils.ValidateIdentifier(arg) {
		return starlark.String(arg), nil
	}

	// Otherwise, error
	return nil, fmt.Errorf("invalid macro argument: %s (not a valid expression or identifier)", arg)
}

// ParseMacroInvocation parses "@name(arg1, arg2, ...)" and returns name and args
func ParseMacroInvocation(invocation string) (string, []string, error) {
	invocation = strings.TrimSpace(invocation)

	// Must start with @
	if !strings.HasPrefix(invocation, "@") {
		return "", nil, fmt.Errorf("macro invocation must start with @")
	}
	rest := invocation[1:]

	// Find opening paren
	parenStart := strings.IndexByte(rest, '(')
	if parenStart == -1 {
		return "", nil, fmt.Errorf("macro invocation missing '(': %s", invocation)
	}

	// Find closing paren
	parenEnd := strings.LastIndexByte(rest, ')')
	if parenEnd == -1 || parenEnd < parenStart {
		return "", nil, fmt.Errorf("macro invocation missing ')': %s", invocation)
	}

	name := strings.TrimSpace(rest[:parenStart])
	if name == "" {
		return "", nil, fmt.Errorf("macro name is empty")
	}

	// Parse arguments (handle nested parens for expressions)
	argStr := strings.TrimSpace(rest[parenStart+1 : parenEnd])
	args, err := splitMacroArgs(argStr)
	if err != nil {
		return "", nil, err
	}

	return name, args, nil
}

// splitMacroArgs splits comma-separated arguments, respecting parentheses
func splitMacroArgs(argStr string) ([]string, error) {
	if argStr == "" {
		return []string{}, nil
	}

	var args []string
	var current strings.Builder
	parenDepth := 0

	for _, ch := range argStr {
		if ch == '(' {
			parenDepth++
			current.WriteRune(ch)
		} else if ch == ')' {
			parenDepth--
			if parenDepth < 0 {
				return nil, fmt.Errorf("unbalanced parentheses in macro arguments")
			}
			current.WriteRune(ch)
		} else if ch == ',' && parenDepth == 0 {
			args = append(args, strings.TrimSpace(current.String()))
			current.Reset()
		} else {
			current.WriteRune(ch)
		}
	}

	if parenDepth != 0 {
		return nil, fmt.Errorf("unbalanced parentheses in macro arguments")
	}

	// Add final argument
	if current.Len() > 0 {
		args = append(args, strings.TrimSpace(current.String()))
	}

	return args, nil
}