c65gm/internal/compiler/scriptexec.go

package compiler

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"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
}

// starlarkErrorMsgLineMatch matches Starlark error messages of the form
// "filename:LINE:COL: message" produced by resolve errors.
var starlarkErrorMsgLineMatch = regexp.MustCompile(`^(.+):(\d+):(\d+): (.*)$`)

// 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)
		}
	}
	// For non-EvalError (resolve errors), extract line from the message
	return parseStarlarkLineFromMsg(err.Error())
}

// parseStarlarkLineFromMsg attempts to extract the Starlark line number from
// a non-EvalError error message string. Starlark resolve errors follow the
// format "filename:LINE:COL: message".
func parseStarlarkLineFromMsg(msg string) int {
	matches := starlarkErrorMsgLineMatch.FindStringSubmatch(msg)
	if len(matches) >= 3 {
		line, err := strconv.Atoi(matches[2])
		if err == nil && line > 0 {
			return line
		}
	}
	return 0
}

// starlarkErrorMsg extracts just the message from a Starlark error.
// For EvalError, it returns the Msg field directly.
// For non-EvalError (resolve errors), it strips the "filename:LINE:COL: " prefix.
func starlarkErrorMsg(err error) string {
	var evalErr *starlark.EvalError
	if errors.As(err, &evalErr) {
		return evalErr.Msg
	}
	return stripStarlarkPositionFromMsg(err.Error())
}

// stripStarlarkPositionFromMsg removes the "filename:LINE:COL: " prefix from
// a Starlark error message if present.
func stripStarlarkPositionFromMsg(msg string) string {
	matches := starlarkErrorMsgLineMatch.FindStringSubmatch(msg)
	if len(matches) >= 5 {
		return matches[4]
	}
	return msg
}

// 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 from pragma or default (prevent infinite loops)
	thread.SetMaxExecutionSteps(readScriptMaxSteps(ctx, scriptLines[0].PragmaSetIndex))

	// Build predeclared: math module + library globals + file I/O builtins
	predeclared := starlark.StringDict{
		"math":        math.Module,
		"load_binary": makeLoadBinary(ctx.ProjectRoot),
		"load_text":   makeLoadText(ctx.ProjectRoot),
	}
	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)
		idx := -1
		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)
			srcLine := scriptLines[idx]
			return nil, fmt.Errorf("Starlark error: %s:%d: %s", srcLine.Filename, srcLine.LineNo, 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
}

// validateScriptFilePath checks that path is safe and resolves it within the project root.
// It rejects absolute paths and path traversal (..).
func validateScriptFilePath(projectRoot, path string) (string, error) {
	if projectRoot == "" {
		return "", fmt.Errorf("project root not set (internal error)")
	}
	if path == "" {
		return "", fmt.Errorf("file path must not be empty")
	}
	if filepath.IsAbs(path) {
		return "", fmt.Errorf("absolute paths are not allowed: %s", path)
	}

	// Reject path traversal components
	cleaned := filepath.Clean(path)
	for _, component := range strings.Split(cleaned, string(filepath.Separator)) {
		if component == ".." {
			return "", fmt.Errorf("path traversal is not allowed: %s", path)
		}
	}

	// Resolve against project root
	resolved := filepath.Join(projectRoot, cleaned)

	// Verify containment within project root
	projectRootWithSep := projectRoot + string(filepath.Separator)
	if !strings.HasPrefix(resolved, projectRootWithSep) && resolved != projectRoot {
		return "", fmt.Errorf("file access denied: path resolves outside project folder")
	}

	// Resolve symlinks to prevent symlink-based escape
	realPath, err := filepath.EvalSymlinks(resolved)
	if err == nil {
		if !strings.HasPrefix(realPath, projectRootWithSep) && realPath != projectRoot {
			return "", fmt.Errorf("file access denied: symlink target resolves outside project folder")
		}
		return realPath, nil
	}

	return resolved, nil
}

// makeLoadBinary creates a Starlark builtin function load_binary(path, offset?, length?)
// that reads a binary file relative to the project root and returns a list of ints (0-255).
func makeLoadBinary(projectRoot string) *starlark.Builtin {
	return starlark.NewBuiltin("load_binary", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var path string
		var offset, length int
		if err := starlark.UnpackPositionalArgs("load_binary", args, kwargs, 1, &path, &offset, &length); err != nil {
			return nil, err
		}

		resolvedPath, err := validateScriptFilePath(projectRoot, path)
		if err != nil {
			return nil, fmt.Errorf("load_binary: %w", err)
		}

		data, err := os.ReadFile(resolvedPath)
		if err != nil {
			return nil, fmt.Errorf("load_binary: cannot read %s: %w", path, err)
		}

		if offset < 0 {
			return nil, fmt.Errorf("load_binary: offset must be non-negative, got %d", offset)
		}
		if offset > len(data) {
			return nil, fmt.Errorf("load_binary: offset %d exceeds file size %d", offset, len(data))
		}
		data = data[offset:]

		if length > 0 {
			if length > len(data) {
				return nil, fmt.Errorf("load_binary: length %d exceeds available data %d", length, len(data))
			}
			data = data[:length]
		}

		result := make([]starlark.Value, len(data))
		for i, b := range data {
			result[i] = starlark.MakeInt(int(b))
		}

		return starlark.NewList(result), nil
	})
}

// makeLoadText creates a Starlark builtin function load_text(path)
// that reads a text file relative to the project root and returns a list of strings (lines).
func makeLoadText(projectRoot string) *starlark.Builtin {
	return starlark.NewBuiltin("load_text", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var path string
		if err := starlark.UnpackPositionalArgs("load_text", args, kwargs, 1, &path); err != nil {
			return nil, err
		}

		resolvedPath, err := validateScriptFilePath(projectRoot, path)
		if err != nil {
			return nil, fmt.Errorf("load_text: %w", err)
		}

		data, err := os.ReadFile(resolvedPath)
		if err != nil {
			return nil, fmt.Errorf("load_text: cannot read %s: %w", path, err)
		}

		text := strings.ReplaceAll(string(data), "\r\n", "\n")
		text = strings.TrimRight(text, "\n")

		var lines []string
		if text == "" {
			lines = []string{}
		} else {
			lines = strings.Split(text, "\n")
		}

		result := make([]starlark.Value, len(lines))
		for i, line := range lines {
			result[i] = starlark.String(line)
		}

		return starlark.NewList(result), nil
	})
}

// readScriptMaxSteps reads the _P_SCRIPT_MAX_STEPS pragma from the given pragma set.
// Returns the configured value (must be > 0), or 1000000 as default.
func readScriptMaxSteps(ctx *CompilerContext, pragmaSetIndex int) uint64 {
	const defaultSteps uint64 = 1000000
	ps := ctx.Pragma.GetPragmaSetByIndex(pragmaSetIndex)
	v := ps.GetPragma("_P_SCRIPT_MAX_STEPS")
	if v == "" {
		return defaultSteps
	}
	n, err := strconv.ParseUint(v, 10, 64)
	if err != nil || n == 0 {
		return defaultSteps
	}
	return n
}

// ExecuteMacro executes a named macro with the given arguments and returns output lines
// pragmaSetIndex is the index of the pragma set at the macro invocation call site.
func ExecuteMacro(macroName string, args []string, ctx *CompilerContext, pragmaSetIndex int) ([]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 from pragma at call site or default
	thread.SetMaxExecutionSteps(readScriptMaxSteps(ctx, pragmaSetIndex))

	// Build predeclared: math + library globals + file I/O builtins + parameter bindings
	predeclared := starlark.StringDict{
		"math":        math.Module,
		"load_binary": makeLoadBinary(ctx.ProjectRoot),
		"load_text":   makeLoadText(ctx.ProjectRoot),
	}
	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
}