ModelRouter/internal/lua/vm.go

// Package lua implements the adapter runtime: bundles/loads *.lua adapter
// scripts, exposes json/string helpers, and lets Go call the protocol
// transform functions plus a build_headers signature hook.
package lua

import (
	"crypto/hmac"
	"crypto/sha256"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	lua "github.com/yuin/gopher-lua"
)

//go:embed adapters/*.lua
var bundledAdapters embed.FS

type APIAdapter struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

type AdapterCache struct {
	mu    sync.RWMutex
	state *lua.LState
	items map[string]*lua.LTable
}

func newAdapterCache() *AdapterCache {
	return &AdapterCache{state: lua.NewState(), items: map[string]*lua.LTable{}}
}

func (c *AdapterCache) setupGlobals() {
	s := c.state
	jsonTable := s.NewTable()
	s.SetGlobal("json", jsonTable)
	s.SetField(jsonTable, "encode", s.NewFunction(func(L *lua.LState) int {
		b, err := jsonEncode(luaValueToGo(L.CheckAny(1)))
		if err != nil {
			L.Push(lua.LString("null"))
			return 1
		}
		L.Push(lua.LString(string(b)))
		return 1
	}))
	s.SetField(jsonTable, "decode", s.NewFunction(func(L *lua.LState) int {
		v, err := jsonDecode(L.CheckString(1))
		if err != nil {
			L.Push(lua.LNil)
			return 1
		}
		L.Push(goValueToLua(L, v))
		return 1
	}))

	// signature/crypto helpers (app verification, timing-safe auth)
	s.SetGlobal("hmac_sha256_hex", s.NewFunction(func(L *lua.LState) int {
		key := L.CheckString(1)
		data := L.CheckString(2)
		m := hmac.New(sha256.New, []byte(key))
		m.Write([]byte(data))
		L.Push(lua.LString(hex.EncodeToString(m.Sum(nil))))
		return 1
	}))
	s.SetGlobal("sha256_hex", s.NewFunction(func(L *lua.LState) int {
		h := sha256.Sum256([]byte(L.CheckString(1)))
		L.Push(lua.LString(hex.EncodeToString(h[:])))
		return 1
	}))
	s.SetGlobal("base64_encode", s.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(base64.StdEncoding.EncodeToString([]byte(L.CheckString(1)))))
		return 1
	}))
	s.SetGlobal("tohex", s.NewFunction(func(L *lua.LState) int {
		L.Push(lua.LString(hex.EncodeToString([]byte(L.CheckString(1)))))
		return 1
	}))

	s.SetGlobal("log", s.NewFunction(func(L *lua.LState) int {
		level := L.ToString(1)
		msg := L.ToString(2)
		fmt.Printf("[adapter/%s] %s\n", level, msg)
		return 0
	}))
}

func (c *AdapterCache) Preload(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("read adapter: %w", err)
	}
	return c.PreloadSource(filepath.Base(path), string(data))
}

func (c *AdapterCache) PreloadSource(name, code string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := c.state.DoString(code); err != nil {
		return fmt.Errorf("compile adapter: %w", err)
	}
	tbl, ok := c.state.Get(-1).(*lua.LTable)
	c.state.Pop(1)
	if !ok {
		return fmt.Errorf("adapter script must return a table")
	}
	if n := tbl.RawGetString("name"); n != nil && n.String() != "" {
		name = n.String()
	}
	c.items[name] = tbl
	return nil
}

func (c *AdapterCache) Get(name string) *lua.LTable {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.items[name]
}

// Remove deletes an adapter from the cache.
func (c *AdapterCache) Remove(name string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.items, name)
}

func (c *AdapterCache) List() []APIAdapter {	c.mu.RLock()
	defer c.mu.RUnlock()
	list := make([]APIAdapter, 0, len(c.items))
	for name, tbl := range c.items {
		a := APIAdapter{Name: name}
		if v := tbl.RawGetString("version"); v != nil {
			a.Version = v.String()
		}
		list = append(list, a)
	}
	return list
}

// VM wraps AdapterCache to dispatch adapter hook calls safely (single Lua
// state is shared, so calls are serialized by a mutex).
type VM struct {
	mu    sync.Mutex
	cache *AdapterCache
	dir   string
}

func NewVM(dir string) *VM {
	return &VM{dir: dir, cache: newAdapterCache()}
}

func (v *VM) Start() error {
	if v.dir != "" {
		if err := os.MkdirAll(v.dir, 0755); err != nil {
			return fmt.Errorf("mkdir adapter dir: %w", err)
		}
		if err := v.writeBundledAdapters(); err != nil {
			return err
		}
		entries, err := os.ReadDir(v.dir)
		if err != nil {
			return err
		}
		for _, e := range entries {
			if filepath.Ext(e.Name()) != ".lua" {
				continue
			}
			if err := v.cache.Preload(filepath.Join(v.dir, e.Name())); err != nil {
				fmt.Printf("[lua] preload %s: %v\n", e.Name(), err)
			}
		}
	}
	v.cache.setupGlobals()
	return nil
}

func (v *VM) Stop() {
	v.mu.Lock()
	defer v.mu.Unlock()
	if v.cache.state != nil {
		v.cache.state.Close()
		v.cache.state = nil
	}
}

func (v *VM) ListAdapters() []APIAdapter { return v.cache.List() }

// LoadAdapter compiles and registers an adapter from a file (runtime safe).
func (v *VM) LoadAdapter(path string) error { return v.cache.Preload(path) }

// RemoveAdapter evicts an adapter from the cache (runtime safe).
func (v *VM) RemoveAdapter(name string) { v.cache.Remove(name) }

func (v *VM) writeBundledAdapters() error {
	known := []string{"openai", "anthropic", "deepseek", "gemini", "github", "groq", "mistral", "ollama", "kimicode"}
	for _, name := range known {
		dst := filepath.Join(v.dir, name+".lua")
		if _, err := os.Stat(dst); err == nil {
			continue
		}
		data, err := bundledAdapters.ReadFile("adapters/" + name + ".lua")
		if err != nil {
			continue
		}
		if err := os.WriteFile(dst, data, 0644); err != nil {
			return err
		}
	}
	return nil
}

func (v *VM) Transform(name, fn, raw string) (string, error) {
	adapter := v.cache.Get(name)
	if adapter == nil {
		return "", fmt.Errorf("adapter %s not loaded", name)
	}
	v.mu.Lock()
	defer v.mu.Unlock()
	f := adapter.RawGetString(fn)
	if f == nil || f == lua.LNil {
		return "", fmt.Errorf("adapter %s missing %s", name, fn)
	}
	state := v.cache.state
	state.Push(f)
	state.Push(lua.LString(raw))
	if err := state.PCall(1, 1, nil); err != nil {
		return "", fmt.Errorf("%s: %w", fn, err)
	}
	res := state.Get(-1)
	state.Pop(1)
	return res.String(), nil
}

// BuildHeaders calls adapter.build_headers(meta). If the adapter does not
// define build_headers, it falls back to the static adapter.headers table.
func (v *VM) BuildHeaders(name string, meta map[string]interface{}) (map[string]string, error) {
	adapter := v.cache.Get(name)
	if adapter == nil {
		return nil, fmt.Errorf("adapter %s not loaded", name)
	}
	v.mu.Lock()
	defer v.mu.Unlock()
	state := v.cache.state

	fn := adapter.RawGetString("build_headers")
	if fn == nil || fn == lua.LNil {
		// fall back to static headers table
		headers := map[string]string{}
		if ht := adapter.RawGetString("headers"); ht != nil {
			if tbl, ok := ht.(*lua.LTable); ok {
				tbl.ForEach(func(key, val lua.LValue) { headers[key.String()] = val.String() })
			}
		}
		return headers, nil
	}
	state.Push(fn)
	state.Push(goValueToLua(state, meta))
	if err := state.PCall(1, 1, nil); err != nil {
		return nil, fmt.Errorf("build_headers: %w", err)
	}
	res := state.Get(-1)
	state.Pop(1)
	headers := map[string]string{}
	if tbl, ok := res.(*lua.LTable); ok {
		tbl.ForEach(func(key, val lua.LValue) {
			k := key.String()
			if k != "" && k != "n" {
				headers[k] = val.String()
			}
		})
	}
	return headers, nil
}

func (v *VM) Endpoint(name string) string {
	adapter := v.cache.Get(name)
	if adapter == nil {
		return ""
	}
	if ep := adapter.RawGetString("endpoint"); ep != nil {
		return ep.String()
	}
	return ""
}

func jsonEncode(v interface{}) ([]byte, error) { return json.Marshal(v) }
func jsonDecode(s string) (interface{}, error) {
	var v interface{}
	if err := json.Unmarshal([]byte(s), &v); err != nil {
		return nil, err
	}
	return v, nil
}

func luaValueToGo(lv lua.LValue) interface{} {
	switch x := lv.(type) {
	case lua.LString:
		return string(x)
	case lua.LNumber:
		return float64(x)
	case lua.LBool:
		return bool(x)
	case *lua.LTable:
		if x.MaxN() > 0 {
			arr := make([]interface{}, 0, x.MaxN())
			x.ForEach(func(_, val lua.LValue) { arr = append(arr, luaValueToGo(val)) })
			return arr
		}
		m := map[string]interface{}{}
		x.ForEach(func(key, val lua.LValue) { m[key.String()] = luaValueToGo(val) })
		return m
	default:
		return nil
	}
}

func goValueToLua(L *lua.LState, val interface{}) lua.LValue {
	switch x := val.(type) {
	case string:
		return lua.LString(x)
	case float64:
		return lua.LNumber(x)
	case int:
		return lua.LNumber(x)
	case int64:
		return lua.LNumber(x)
	case bool:
		return lua.LBool(x)
	case nil:
		return lua.LNil
	case []interface{}:
		t := L.NewTable()
		for i, item := range x {
			t.RawSetInt(i+1, goValueToLua(L, item))
		}
		return t
	case map[string]interface{}:
		t := L.NewTable()
		for k, item := range x {
			t.RawSetString(k, goValueToLua(L, item))
		}
		return t
	default:
		return lua.LNil
	}
}