gojs/gojs.go

package gojs

import (
	"apigo.cc/apigo/gojs/dop251/goja"
	"apigo.cc/apigo/gojs/dop251/goja_nodejs/require"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/ssgo/log"
	"github.com/ssgo/tool/watcher"
	"github.com/ssgo/u"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"
)

type Map = map[string]any

type Module struct {
	Object      Map
	ObjectMaker func(vm *goja.Runtime) Map
	TsCode      string
	Desc        string
	Example     string
}

type Program struct {
	prg       *goja.Program
	startFile string
	imports   map[string]string
}

var modules = map[string]Module{}
var modulesLock = sync.RWMutex{}

func Register(name string, mod Module) {
	modulesLock.Lock()
	modules[name] = mod
	modulesLock.Unlock()
}

func RunFile(file string, args ...any) (any, error) {
	if code, err := u.ReadFile(file); err == nil {
		return Run(code, file, args...)
	} else {
		return nil, err
	}
}

func Run(code string, refFile string, args ...any) (any, error) {
	rt := New()
	var r any
	err := rt.StartFromCode(code, refFile)
	if err == nil {
		r, err = rt.RunMain(args...)
	}
	return r, err
}

func RunProgram(plg *Program, args ...any) (any, error) {
	rt := New()
	var r any
	err := rt.StartFromProgram(plg)
	if err == nil {
		r, err = rt.RunMain(args...)
	}
	return r, err
}

func CompileFile(file string) (*Program, error) {
	if code, err := u.ReadFile(file); err == nil {
		return CompileMain(code, file)
	} else {
		return nil, err
	}
}

func CompileMain(code string, refFile string) (*Program, error) {
	imports := makeImports(&code)
	if !checkMainMatcher.MatchString(code) {
		code = "function main(...Args){" + code + "}"
	}
	p, err := goja.Compile(refFile, code, false)
	return &Program{prg: p, imports: imports, startFile: refFile}, err
}

func CompileCode(code string, refFile string) (*Program, error) {
	imports := makeImports(&code)
	p, err := goja.Compile(refFile, code, false)
	return &Program{prg: p, imports: imports, startFile: refFile}, err
}

var importModMatcher = regexp.MustCompile(`(?im)^\s*import\s+(.+?)\s+from\s+['"](.+?)['"]`)
var requireModMatcher = regexp.MustCompile(`(?im)^\s*(const|let|var)\s+(.+?)\s*=\s*require\s*\(\s*['"](.+?)['"]\s*\)`)

// var importLibMatcher = regexp.MustCompile(`(?im)^\s*(import)\s+(.+?)\s+from\s+['"][./\\\w:]+lib[/\\](.+?)(\.ts)?['"]`)
// var requireLibMatcher = regexp.MustCompile(`(?im)^\s*(const|let|var)\s+(.+?)\s*=\s*require\s*\(\s*['"][./\\\w:]+lib[/\\](.+?)(\.ts)?['"]\s*\)`)
var checkMainMatcher = regexp.MustCompile(`(?im)^\s*function\s+main\s*\(`)

type Runtime struct {
	vm           *goja.Runtime
	required     map[string]bool
	file         string
	srcCode      string
	code         string
	moduleLoader func(string) string
	started      bool
	dataLock     sync.RWMutex
}

func (rt *Runtime) lock() {
	rt.vm.Locker.Lock()
}

func (rt *Runtime) unlock() {
	rt.vm.Locker.Unlock()
}

//func (rt *Runtime) Free() {
//	rt.vm.GoData = nil
//	rt.vm = nil
//}

func (rt *Runtime) SetModuleLoader(fn func(filename string) string) {
	rt.moduleLoader = fn
}

func (rt *Runtime) GetCallStack() []string {
	callStacks := make([]string, 0)
	rt.lock()
	stacks := rt.vm.CaptureCallStack(0, nil)
	rt.unlock()
	for _, stack := range stacks {
		callStacks = append(callStacks, stack.Position().String())
	}
	return callStacks
}

func (rt *Runtime) requireSpecialMod(name string) error {
	if rt.required[name] {
		return nil
	}
	if name == "setTimeout" {
		rt.required["setTimeout"] = true
		fn := func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := MakeArgs(&argsIn, vm).Check(1)
			callback := args.Func(0)
			timeout := time.Duration(args.Int64(1)) * time.Millisecond
			//locked := vm.Locker.TryLock()
			//vm.Locker.Unlock()
			//locked2 := vm.CallbackLocker.TryLock()
			//vm.CallbackLocker.Unlock()
			if callback != nil {
				go func() {
					if timeout > 0 {
						time.Sleep(timeout)
					}
					//if !locked {
					//	vm.Locker.Lock()
					//}
					//if !locked2 {
					//	vm.CallbackLocker.Lock()
					//}
					if _, err := callback(args.This, args.Arguments[2:]...); err != nil {
						//panic(vm.NewGoError(err))
					}
				}()
			}
			return nil
		}
		rt.lock()
		err := rt.vm.Set("setTimeout", fn)
		rt.unlock()
		return err
	}

	return errors.New("module not found: " + name)
}

func (rt *Runtime) requireMod(modName, realModName string) error {
	if rt.required[modName] {
		return nil
	}
	modulesLock.RLock()
	mod, ok := modules[modName]
	modulesLock.RUnlock()
	if !ok {
		// 不在注册的模块中，尝试引入特殊模块
		return rt.requireSpecialMod(modName)
	}

	var err error
	rt.lock()
	if mod.ObjectMaker != nil {
		err = rt.vm.Set(realModName, mod.ObjectMaker(rt.vm))
	} else {
		err = rt.vm.Set(realModName, mod.Object)
	}
	rt.unlock()
	if err != nil {
		return err
	}
	rt.required[modName] = true
	return nil
}

func (rt *Runtime) requireAllMod() error {
	var allModules = map[string]Module{}
	modulesLock.RLock()
	for k, v := range modules {
		allModules[k] = v
	}
	modulesLock.RUnlock()

	for modName, mod := range allModules {
		if !rt.required[modName] {
			var err error
			realModName := modName
			if strings.ContainsRune(modName, '/') {
				realModName = strings.ReplaceAll(realModName, "/", "_")
			}
			rt.lock()
			if mod.ObjectMaker != nil {
				err = rt.vm.Set(realModName, mod.ObjectMaker(rt.vm))
			} else {
				err = rt.vm.Set(realModName, mod.Object)
			}
			rt.unlock()
			if err != nil {
				return err
			}
			rt.required[modName] = true
		}
	}
	if err := rt.requireSpecialMod("setTimeout"); err != nil {
		return err
	}
	return nil
}

func makeImports(code *string) map[string]string {
	importModules := map[string]string{}
	*code = importModMatcher.ReplaceAllString(*code, "let $1 = require('$2')")
	*code = requireModMatcher.ReplaceAllStringFunc(*code, func(str string) string {
		if m := requireModMatcher.FindStringSubmatch(str); m != nil && len(m) > 3 {
			optName := m[1]
			varName := m[2]
			modName := m[3]
			realModName := modName
			if strings.ContainsRune(modName, '/') {
				realModName = strings.ReplaceAll(realModName, "/", "_")
			}
			modulesLock.RLock()
			_, ok := modules[modName]
			modulesLock.RUnlock()
			if !ok {
				return str
			}
			importModules[modName] = realModName
			if varName != realModName {
				return fmt.Sprintf("%s %s = %s", optName, varName, realModName)
			} else {
				return ""
			}
		}
		return str
	})

	if strings.Contains(*code, "setTimeout") {
		importModules["setTimeout"] = "setTimeout"
	}
	return importModules
}

func (rt *Runtime) setImports(importModules map[string]string) error {
	var modErr error
	for k, v := range importModules {
		if modErr == nil {
			if err := rt.requireMod(k, v); err != nil {
				modErr = err
			}
		}
	}
	return modErr
}

func New() *Runtime {
	vm := goja.New()

	rt := &Runtime{
		vm:       vm,
		required: map[string]bool{},
	}

	vm.GoData = map[string]any{
		"logger": log.New(u.ShortUniqueId()),
	}

	// 处理模块引用
	require.NewRegistryWithLoader(func(path string) ([]byte, error) {
		modFile := path
		if !filepath.IsAbs(modFile) {
			modFile = filepath.Join(filepath.Dir(rt.file), modFile)
		}
		if !strings.HasSuffix(modFile, ".js") && !u.FileExists(modFile) {
			modFile += ".js"
		}
		modCode := ""
		if rt.moduleLoader != nil {
			modCode = rt.moduleLoader(modFile)
		}
		if modCode == "" {
			var err error
			modCode, err = u.ReadFile(modFile)
			if err != nil {
				return nil, err
			}
		}
		imports := makeImports(&modCode)
		if err := rt.setImports(imports); err != nil {
			return nil, err
		}
		return []byte(modCode), nil
	}).Enable(rt.vm)

	return rt
}

func (rt *Runtime) StartFromFile(file string) error {
	if code, err := u.ReadFile(file); err == nil {
		return rt.StartFromCode(code, file)
	} else {
		return err
	}
}

func (rt *Runtime) StartFromCode(code, refFile string) error {
	if refFile != "" {
		rt.file = refFile
	}
	if rt.file == "" {
		rt.file = "main.js"
	}

	if absFile, err := filepath.Abs(rt.file); err == nil {
		rt.file = absFile
	}
	refPath := filepath.Dir(refFile)
	rt.SetGoData("startFile", refFile)
	rt.SetGoData("startPath", refPath)

	if rt.srcCode == "" {
		rt.srcCode = code
	}
	rt.code = code

	// 按需加载引用
	//var importCount int
	var modErr error
	//rt.code, importCount, modErr = rt.makeImport(rt.code)
	imports := makeImports(&rt.code)
	if len(imports) == 0 {
		modErr = rt.requireAllMod()
	} else {
		modErr = rt.setImports(imports)
	}

	//// 如果没有import，默认import所有
	//if modErr == nil && importCount == 0 {
	//	modErr = rt.requireMod("", "")
	//}
	if modErr != nil {
		return modErr
	}

	//fmt.Println(u.BCyan(rt.code))

	// 初始化主函数
	if !checkMainMatcher.MatchString(rt.code) {
		rt.code = "function main(...Args){" + rt.code + "}"
	}
	rt.lock()
	_, err := rt.vm.RunScript(rt.file, rt.code)
	rt.unlock()
	if err != nil {
		return err
	} else {
		rt.started = true
		return nil
	}
}

func (rt *Runtime) StartFromProgram(plg *Program) error {
	var modErr error
	if len(plg.imports) == 0 {
		modErr = rt.requireAllMod()
	} else {
		modErr = rt.setImports(plg.imports)
	}
	if modErr != nil {
		return modErr
	}
	rt.SetGoData("startFile", plg.startFile)
	rt.SetGoData("startPath", filepath.Dir(plg.startFile))
	rt.lock()
	_, err := rt.vm.RunProgram(plg.prg)
	rt.unlock()
	if err != nil {
		return err
	} else {
		rt.started = true
		return nil
	}
}

func (rt *Runtime) RunMain(args ...any) (any, error) {
	if !rt.started {
		return nil, errors.New("runtime not started")
	}

	// 解析参数
	for i, arg := range args {
		if str, ok := arg.(string); ok {
			var v interface{}
			if err := json.Unmarshal([]byte(str), &v); err == nil {
				args[i] = v
			}
		}
	}

	rt.lock()
	err := rt.vm.Set("__args", args)
	rt.unlock()
	if err != nil {
		return nil, err
	}
	rt.lock()
	r, err := rt.vm.RunScript("main", "main(...__args)")
	rt.unlock()
	return makeResult(r, err)
}

func (rt *Runtime) RunVM(callback func(vm *goja.Runtime) (any, error)) (any, error) {
	rt.lock()
	r, err := callback(rt.vm)
	rt.unlock()
	return r, err
}

func (rt *Runtime) SetGoData(name string, value any) {
	rt.dataLock.Lock()
	rt.vm.GoData[name] = value
	rt.dataLock.Unlock()
}

func (rt *Runtime) GetGoData(name string) any {
	rt.dataLock.RLock()
	defer rt.dataLock.RUnlock()
	return rt.vm.GoData[name]
}

func (rt *Runtime) Set(name string, value any) error {
	rt.lock()
	defer rt.unlock()
	return rt.vm.Set(name, value)
}

func (rt *Runtime) SetGlobal(global Map) error {
	var err error
	rt.lock()
	defer rt.unlock()
	for k, v := range global {
		if err = rt.vm.Set(k, v); err != nil {
			return err
		}
	}
	return nil
}

func (rt *Runtime) RunFile(file string) (any, error) {
	if code, err := u.ReadFile(file); err == nil {
		imports := makeImports(&code)
		if err := rt.setImports(imports); err != nil {
			return nil, err
		}
		rt.lock()
		r, err := rt.vm.RunScript(file, code)
		rt.unlock()
		return makeResult(r, err)
	} else {
		return nil, err
	}
}

func (rt *Runtime) RunCode(code string) (any, error) {
	imports := makeImports(&code)
	if err := rt.setImports(imports); err != nil {
		return nil, err
	}
	rt.lock()
	r, err := rt.vm.RunScript("anonymous", code)
	rt.unlock()
	return makeResult(r, err)
}

func (rt *Runtime) RunProgram(prg *Program) (any, error) {
	var r goja.Value
	err := rt.setImports(prg.imports)
	if err == nil {
		rt.lock()
		r, err = rt.vm.RunProgram(prg.prg)
		rt.unlock()
	}
	return makeResult(r, err)
}

func makeResult(r goja.Value, err error) (any, error) {
	var result any
	if err == nil {
		if r != nil && !r.Equals(goja.Undefined()) {
			result = r.Export()
		}
	}
	return result, err
}

type WatchRunner struct {
	w *watcher.Watcher
}

func (wr *WatchRunner) WaitForKill() {
	exitCh := make(chan os.Signal, 1)
	closeCh := make(chan bool, 1)
	signal.Notify(exitCh, os.Interrupt, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP)
	go func() {
		<-exitCh
		closeCh <- true
	}()
	<-closeCh
	wr.w.Stop()
}

func (wr *WatchRunner) Stop() {
	wr.w.Stop()
}

func WatchRun(file string, extDirs, extTypes []string, args ...any) (*WatchRunner, error) {
	wr := &WatchRunner{}
	run := func() {
		rt := New()
		if wr.w != nil {
			rt.SetModuleLoader(func(filename string) string {
				filePath := filepath.Dir(filename)
				needWatch := true
				for _, v := range wr.w.WatchList() {
					if v == filePath {
						needWatch = false
						break
					}
				}
				if needWatch {
					fmt.Println(u.BMagenta("[watching module path]"), filePath)
					_ = wr.w.Add(filePath)
				}
				return u.ReadFileN(filename)
			})
		}
		err := rt.StartFromFile(file)
		result, err := rt.RunMain(args...)
		if err != nil {
			fmt.Println(u.BRed(err.Error()))
			fmt.Println(u.Red("  " + strings.Join(rt.GetCallStack(), "\n  ")))
		} else if result != nil {
			fmt.Println(u.Cyan(u.JsonP(result)))
		}
	}

	var isWaitingRun = false
	onChange := func(filename string, event string) {
		if !isWaitingRun {
			_, _ = os.Stdout.WriteString("\x1b[3;J\x1b[H\x1b[2J")
			isWaitingRun = true
			go func() {
				time.Sleep(time.Millisecond * 10)
				isWaitingRun = false
				run()
			}()
		}
		fmt.Println(u.BYellow("[changed]"), filename)
	}
	_, _ = os.Stdout.WriteString("\x1b[3;J\x1b[H\x1b[2J")
	watchStartPath := filepath.Dir(file)
	fmt.Println(u.BMagenta("[watching root path]"), watchStartPath)
	watchDirs := []string{watchStartPath}
	watchTypes := []string{"js", "json", "yml"}
	if extDirs != nil {
		for _, v := range extDirs {
			watchDirs = append(watchDirs, v)
		}
	}
	if extTypes != nil {
		for _, v := range extTypes {
			watchTypes = append(watchTypes, v)
		}
	}
	if w, err := watcher.Start(watchDirs, watchTypes, onChange); err == nil {
		wr.w = w
		go func() {
			run()
		}()
		return wr, nil
	} else {
		return nil, err
	}
}

func ExportForDev() string {
	var allModules = map[string]Module{}
	modulesLock.RLock()
	for k, v := range modules {
		allModules[k] = v
	}
	modulesLock.RUnlock()

	dir, _ := os.Getwd()
	nodeModulesPath := "node_modules"
	packageJsonFile := "package.json"
	for i := 0; i < 20; i++ {
		nodePath := filepath.Join(dir, "node_modules")
		if u.FileExists(nodePath) {
			nodeModulesPath = nodePath
			packageJsonFile = filepath.Join(dir, "package.json")
			break
		}
		parentDir := filepath.Dir(dir)
		if parentDir == dir {
			break
		}
		dir = parentDir
	}

	oldPackageJson := u.ReadFileN(packageJsonFile)
	insertModulesPostfix := ""
	if oldPackageJson == "" {
		oldPackageJson = "{\n  \"devDependencies\": {\n\n  }\n}"
	} else if !strings.Contains(oldPackageJson, "\"devDependencies\": {") {
		oldPackageJson = strings.Replace(oldPackageJson, "{", "{\n  \"devDependencies\": {\n\n  },", 1)
	} else {
		insertModulesPostfix = ","
	}

	insertModules := make([]string, 0)
	imports := make([]string, len(allModules))
	i := 0
	for k, v := range allModules {
		varName := k
		tsPath := k
		if strings.ContainsRune(k, '/') {
			varName = varName[strings.LastIndex(varName, "/")+1:]
			if os.PathSeparator != '/' {
				tsPath = strings.ReplaceAll(tsPath, "/", string(os.PathSeparator))
			}
		}
		_ = u.WriteFile(filepath.Join(nodeModulesPath, tsPath, "index.ts"), v.TsCode)
		imports[i] = fmt.Sprintf("import %s from '%s'", varName, k)
		i++

		if !strings.Contains(oldPackageJson, "\""+k+"\":") {
			insertModules = u.AppendUniqueString(insertModules, fmt.Sprint("\"", k, "\": \"v0.0.0\""))
		}
	}

	if len(insertModules) > 0 {
		newPackageJson := strings.Replace(oldPackageJson, "\"devDependencies\": {", "\"devDependencies\": {\n    "+strings.Join(insertModules, ",\n    ")+insertModulesPostfix, 1)
		_ = u.WriteFile(packageJsonFile, newPackageJson)
	}
	return strings.Join(imports, "\n")
}