package compiler

import (
	"bytes"
	"fmt"
	"strings"

	"c65gm/internal/utils"

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

// 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 []string, ctx *CompilerContext, isLibrary bool) ([]string, error) {
	// Join script lines
	scriptText := strings.Join(scriptLines, "\n")

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

	var finalScript string
	if isLibrary {
		// LIBRARY: execute at top level so defs become globals
		finalScript = scriptText
	} 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"
	}

	// 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, "script.star", finalScript, predeclared)
	if err != nil {
		return nil, 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"

	// 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, "macro.star", finalScript, predeclared)
	if err != nil {
		return nil, 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 := func(name string) (int64, bool) {
		sym := ctx.SymbolTable.Lookup(name, nil)
		if sym != nil && sym.IsConst() {
			return int64(sym.Value), true
		}
		return 0, false
	}

	// 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
}