Moonshark/runner/templates.go

package runner

import (
	"fmt"
	"html"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"maps"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// CachedTemplate holds compiled Lua code with metadata
type CachedTemplate struct {
	CompiledLua []byte // Compiled Lua bytecode
	ModTime     time.Time
	Path        string
}

// TemplateCache manages template caching with fast lookups
type TemplateCache struct {
	templates sync.Map // map[string]*CachedTemplate
}

var (
	templateCache = &TemplateCache{}
	simpleVarRe   = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
)

// htmlEscape escapes HTML special characters
func htmlEscape(s string) string {
	return html.EscapeString(s)
}

// compileTemplate converts template string to Lua code
func compileTemplate(templateStr string) (string, error) {
	pos := 1
	chunks := []any{}
	templateLen := len(templateStr)

	for pos <= templateLen {
		// Find next template tag
		unescapedStart := strings.Index(templateStr[pos-1:], "{{{")
		escapedStart := strings.Index(templateStr[pos-1:], "{{")

		// Adjust positions to be absolute
		if unescapedStart != -1 {
			unescapedStart += pos - 1
		}
		if escapedStart != -1 {
			escapedStart += pos - 1
		}

		var start, openLen int
		var tagType string

		// Determine which tag comes first
		if unescapedStart != -1 && (escapedStart == -1 || unescapedStart <= escapedStart) {
			start, tagType, openLen = unescapedStart, "-", 3
		} else if escapedStart != -1 {
			start, tagType, openLen = escapedStart, "=", 2
		} else {
			// No more tags, add remaining text
			if pos <= templateLen {
				chunks = append(chunks, templateStr[pos-1:])
			}
			break
		}

		// Add text before tag
		if start > pos-1 {
			chunks = append(chunks, templateStr[pos-1:start])
		}

		// Find closing tag
		pos = start + openLen + 1
		var closeTag string
		if tagType == "-" {
			closeTag = "}}}"
		} else {
			closeTag = "}}"
		}

		closeStart := strings.Index(templateStr[pos-1:], closeTag)
		if closeStart == -1 {
			return "", fmt.Errorf("failed to find closing tag at position %d", pos)
		}
		closeStart += pos - 1

		// Extract and trim code
		code := strings.TrimSpace(templateStr[pos-1 : closeStart])

		// Check if it's a simple variable for escaped output
		isSimpleVar := tagType == "=" && simpleVarRe.MatchString(code)

		chunks = append(chunks, []any{tagType, code, pos, isSimpleVar})

		// Move past closing tag
		pos = closeStart + len(closeTag) + 1
	}

	// Generate Lua code
	buffer := []string{"local _tostring, _escape, _b, _b_i = ...\n"}

	for _, chunk := range chunks {
		switch v := chunk.(type) {
		case string:
			// Literal text
			buffer = append(buffer, "_b_i = _b_i + 1\n")
			buffer = append(buffer, fmt.Sprintf("_b[_b_i] = %q\n", v))
		case []any:
			tagType := v[0].(string)
			code := v[1].(string)
			pos := v[2].(int)
			isSimpleVar := v[3].(bool)

			switch tagType {
			case "=":
				if isSimpleVar {
					buffer = append(buffer, "_b_i = _b_i + 1\n")
					buffer = append(buffer, fmt.Sprintf("--[[%d]] _b[_b_i] = _escape(_tostring(%s))\n", pos, code))
				} else {
					buffer = append(buffer, fmt.Sprintf("--[[%d]] %s\n", pos, code))
				}
			case "-":
				buffer = append(buffer, "_b_i = _b_i + 1\n")
				buffer = append(buffer, fmt.Sprintf("--[[%d]] _b[_b_i] = _tostring(%s)\n", pos, code))
			}
		}
	}

	buffer = append(buffer, "return _b")
	return strings.Join(buffer, ""), nil
}

// getTemplate loads or retrieves a cached template
func getTemplate(templatePath string, state *luajit.State) ([]byte, error) {
	// Resolve the path using the fs system
	fullPath, err := ResolvePath(templatePath)
	if err != nil {
		return nil, fmt.Errorf("template path resolution failed: %w", err)
	}

	// Check if file exists and get mod time
	info, err := os.Stat(fullPath)
	if err != nil {
		return nil, fmt.Errorf("template file not found: %w", err)
	}

	// Fast path: check cache first
	if cached, ok := templateCache.templates.Load(templatePath); ok {
		cachedTpl := cached.(*CachedTemplate)
		// Compare mod times for cache validity
		if cachedTpl.ModTime.Equal(info.ModTime()) {
			return cachedTpl.CompiledLua, nil
		}
	}

	// Cache miss or file changed - load and compile template
	content, err := os.ReadFile(fullPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read template: %w", err)
	}

	// Compile template to Lua code
	luaCode, err := compileTemplate(string(content))
	if err != nil {
		return nil, fmt.Errorf("template compile error: %w", err)
	}

	// Compile Lua code to bytecode
	bytecode, err := state.CompileBytecode(luaCode, templatePath)
	if err != nil {
		return nil, fmt.Errorf("lua compile error: %w", err)
	}

	// Store in cache
	cachedTpl := &CachedTemplate{
		CompiledLua: bytecode,
		ModTime:     info.ModTime(),
		Path:        fullPath,
	}
	templateCache.templates.Store(templatePath, cachedTpl)

	return bytecode, nil
}

// goHtmlEscape provides HTML escaping from Go
func goHtmlEscape(state *luajit.State) int {
	if !state.IsString(1) {
		state.PushString("")
		return 1
	}

	input := state.ToString(1)
	escaped := htmlEscape(input)
	state.PushString(escaped)
	return 1
}

// templateInclude renders a template with auto-merged data
func templateInclude(state *luajit.State) int {
	// Get template path
	if !state.IsString(1) {
		state.PushString("template.include: path must be a string")
		return 1
	}
	templatePath := state.ToString(1)

	// Get current template data
	state.GetGlobal("__template_data")
	currentData, _ := state.ToTable(-1)
	state.Pop(1)

	// Get new data (optional)
	var newData map[string]any
	if state.GetTop() >= 2 && !state.IsNil(2) {
		if envValue, err := state.ToValue(2); err == nil {
			if envMap, ok := envValue.(map[string]any); ok {
				newData = envMap
			}
		}
	}

	// Merge data
	mergedData := make(map[string]any)
	if currentData != nil {
		maps.Copy(mergedData, currentData)
	}
	if newData != nil {
		maps.Copy(mergedData, newData)
	}

	// Call templateRender with merged data
	state.PushString(templatePath)
	state.PushTable(mergedData)
	return templateRender(state)
}
func templateRender(state *luajit.State) int {
	// Get template path (required)
	if !state.IsString(1) {
		state.PushString("template.render: template path must be a string")
		return 1
	}
	templatePath := state.ToString(1)

	// Get data (optional)
	var env map[string]any
	if state.GetTop() >= 2 && !state.IsNil(2) {
		var err error
		envValue, err := state.ToValue(2)
		if err != nil {
			state.PushString("template.render: invalid data: " + err.Error())
			return 1
		}
		if envMap, ok := envValue.(map[string]any); ok {
			env = envMap
		}
	}

	// Load compiled template from cache
	bytecode, err := getTemplate(templatePath, state)
	if err != nil {
		state.PushString("template.render: " + err.Error())
		return 1
	}

	// Load bytecode
	if err := state.LoadBytecode(bytecode, templatePath); err != nil {
		state.PushString("template.render: load error: " + err.Error())
		return 1
	}

	// Create runtime environment
	runtimeEnv := make(map[string]any)
	if env != nil {
		maps.Copy(runtimeEnv, env)
	}

	// Add current template data for nested calls
	runtimeEnv["__template_data"] = env

	// Get current global environment
	state.GetGlobal("_G")
	globalEnv, err := state.ToTable(-1)
	if err == nil {
		for k, v := range globalEnv {
			if _, exists := runtimeEnv[k]; !exists {
				runtimeEnv[k] = v
			}
		}
	}
	state.Pop(1)

	// Set up runtime environment
	if err := state.PushTable(runtimeEnv); err != nil {
		state.Pop(1) // Pop bytecode
		state.PushString("template.render: env error: " + err.Error())
		return 1
	}

	// Create metatable for environment
	state.NewTable()
	state.GetGlobal("_G")
	state.SetField(-2, "__index")
	state.SetMetatable(-2)

	// Set environment using setfenv
	state.GetGlobal("setfenv")
	state.PushCopy(-3) // Template function
	state.PushCopy(-3) // Environment
	state.Call(2, 1)   // setfenv(fn, env) returns the function

	// Prepare arguments for template execution
	state.GetGlobal("tostring")        // tostring function
	state.PushGoFunction(goHtmlEscape) // HTML escape function
	state.NewTable()                   // output buffer
	state.PushNumber(0)                // buffer index

	// Execute template (4 args, 1 result)
	if err := state.Call(4, 1); err != nil {
		state.Pop(1) // Pop environment
		state.PushString("template.render: execution error: " + err.Error())
		return 1
	}

	// Get result buffer
	buffer, err := state.ToTable(-1)
	if err != nil {
		state.Pop(2) // Pop buffer and environment
		state.PushString("template.render: result error: " + err.Error())
		return 1
	}

	// Convert buffer to string
	var result strings.Builder
	i := 1
	for {
		if val, exists := buffer[fmt.Sprintf("%d", i)]; exists {
			result.WriteString(fmt.Sprintf("%v", val))
			i++
		} else {
			break
		}
	}

	state.Pop(2) // Pop buffer and environment
	state.PushString(result.String())
	return 1
}

// templateExists checks if a template file exists
func templateExists(state *luajit.State) int {
	if !state.IsString(1) {
		state.PushBoolean(false)
		return 1
	}
	templatePath := state.ToString(1)

	// Resolve path
	fullPath, err := ResolvePath(templatePath)
	if err != nil {
		state.PushBoolean(false)
		return 1
	}

	// Check if file exists
	_, err = os.Stat(fullPath)
	state.PushBoolean(err == nil)
	return 1
}

// templateClearCache clears the template cache
func templateClearCache(state *luajit.State) int {
	// Optional: clear specific template
	if state.GetTop() >= 1 && state.IsString(1) {
		templatePath := state.ToString(1)
		templateCache.templates.Delete(templatePath)
		state.PushBoolean(true)
		return 1
	}

	// Clear entire cache
	templateCache.templates.Range(func(key, value any) bool {
		templateCache.templates.Delete(key)
		return true
	})

	state.PushBoolean(true)
	return 1
}

// templateCacheSize returns the number of templates in cache
func templateCacheSize(state *luajit.State) int {
	count := 0
	templateCache.templates.Range(func(key, value any) bool {
		count++
		return true
	})

	state.PushNumber(float64(count))
	return 1
}

// RegisterTemplateFunctions registers template functions with the Lua state
func RegisterTemplateFunctions(state *luajit.State) error {
	if err := state.RegisterGoFunction("__template_render", templateRender); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__template_include", templateInclude); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__template_exists", templateExists); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__template_clear_cache", templateClearCache); err != nil {
		return err
	}
	if err := state.RegisterGoFunction("__template_cache_size", templateCacheSize); err != nil {
		return err
	}
	return nil
}

// CleanupTemplate clears the template cache
func CleanupTemplate() {
	templateCache.templates.Range(func(key, value any) bool {
		templateCache.templates.Delete(key)
		return true
	})
}

// InvalidateTemplate removes a specific template from cache
func InvalidateTemplate(templatePath string) {
	templateCache.templates.Delete(templatePath)
}