Moonshark/modules/http/http.go

package http

import (
	"context"
	"fmt"
	"net"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/valyala/fasthttp"
)

var (
	globalServer       *fasthttp.Server
	globalWorkerPool   *WorkerPool
	globalStateCreator StateCreator
	globalMu           sync.RWMutex
	serverRunning      bool
	staticHandlers     = make(map[string]*fasthttp.FS)
	staticMu           sync.RWMutex
)

func SetStateCreator(creator StateCreator) {
	globalStateCreator = creator
}

func GetFunctionList() map[string]luajit.GoFunction {
	return map[string]luajit.GoFunction{
		"http_create_server":   http_create_server,
		"http_spawn_workers":   http_spawn_workers,
		"http_listen":          http_listen,
		"http_close_server":    http_close_server,
		"http_has_servers":     http_has_servers,
		"http_register_static": http_register_static,
	}
}

func http_create_server(s *luajit.State) int {
	globalMu.Lock()
	defer globalMu.Unlock()

	if globalServer != nil {
		s.PushBoolean(true) // Already created
		return 1
	}

	globalServer = &fasthttp.Server{
		Handler:      handleRequest,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	s.PushBoolean(true)
	return 1
}

func http_spawn_workers(s *luajit.State) int {
	globalMu.Lock()
	defer globalMu.Unlock()

	if globalWorkerPool != nil {
		s.PushBoolean(true) // Already spawned
		return 1
	}

	if globalStateCreator == nil {
		s.PushBoolean(false)
		s.PushString("state creator not set")
		return 2
	}

	workerCount := max(runtime.NumCPU(), 2)

	pool, err := NewWorkerPool(workerCount, s, globalStateCreator)
	if err != nil {
		s.PushBoolean(false)
		s.PushString(fmt.Sprintf("failed to create worker pool: %v", err))
		return 2
	}
	globalWorkerPool = pool

	s.PushBoolean(true)
	return 1
}

func http_listen(s *luajit.State) int {
	if err := s.CheckMinArgs(1); err != nil {
		return s.PushError("http_listen: %v", err)
	}

	addr := s.ToString(1)

	globalMu.RLock()
	server := globalServer
	globalMu.RUnlock()

	if server == nil {
		s.PushBoolean(false)
		s.PushString("no server created")
		return 2
	}

	globalMu.Lock()
	if serverRunning {
		globalMu.Unlock()
		s.PushBoolean(true) // Already running
		return 1
	}
	serverRunning = true
	globalMu.Unlock()

	go func() {
		if err := server.ListenAndServe(addr); err != nil {
			fmt.Printf("HTTP server error: %v\n", err)
		}
	}()

	time.Sleep(100 * time.Millisecond)

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		globalMu.Lock()
		serverRunning = false
		globalMu.Unlock()
		s.PushBoolean(false)
		s.PushString(fmt.Sprintf("failed to start server: %v", err))
		return 2
	}
	conn.Close()

	s.PushBoolean(true)
	return 1
}

func http_close_server(s *luajit.State) int {
	globalMu.Lock()
	defer globalMu.Unlock()

	if globalWorkerPool != nil {
		globalWorkerPool.Close()
		globalWorkerPool = nil
	}

	if globalServer != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		globalServer.ShutdownWithContext(ctx)
		globalServer = nil
	}

	serverRunning = false
	s.PushBoolean(true)
	return 1
}

func http_has_servers(s *luajit.State) int {
	globalMu.RLock()
	running := serverRunning
	globalMu.RUnlock()

	s.PushBoolean(running)
	return 1
}

func http_register_static(s *luajit.State) int {
	if err := s.CheckMinArgs(2); err != nil {
		return s.PushError("http_register_static: %v", err)
	}

	urlPrefix := s.ToString(1)
	rootPath := s.ToString(2)

	// Ensure prefix starts with /
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}

	// Convert to absolute path
	absPath, err := filepath.Abs(rootPath)
	if err != nil {
		s.PushBoolean(false)
		s.PushString(fmt.Sprintf("invalid path: %v", err))
		return 2
	}

	RegisterStaticHandler(urlPrefix, absPath)
	s.PushBoolean(true)
	return 1
}

func HasActiveServers() bool {
	globalMu.RLock()
	defer globalMu.RUnlock()
	return serverRunning
}

func WaitForServers() {
	for HasActiveServers() {
		time.Sleep(100 * time.Millisecond)
	}
}

func handleRequest(ctx *fasthttp.RequestCtx) {
	path := string(ctx.Path())

	// Check static handlers first
	staticMu.RLock()
	for prefix, fs := range staticHandlers {
		if strings.HasPrefix(path, prefix) {
			staticMu.RUnlock()
			// Remove prefix and serve
			ctx.Request.URI().SetPath(strings.TrimPrefix(path, prefix))
			fs.NewRequestHandler()(ctx)
			return
		}
	}
	staticMu.RUnlock()

	// Fall back to Lua handling
	globalMu.RLock()
	pool := globalWorkerPool
	globalMu.RUnlock()

	if pool == nil {
		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
		ctx.SetBodyString("Worker pool not initialized")
		return
	}

	worker := pool.Get()
	defer pool.Put(worker)

	if worker == nil {
		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
		ctx.SetBodyString("No worker available")
		return
	}

	req := GetRequest()
	defer PutRequest(req)

	resp := GetResponse()
	defer PutResponse(resp)

	// Populate request
	req.Method = string(ctx.Method())
	req.Path = string(ctx.Path())
	req.Body = string(ctx.Request.Body())

	ctx.QueryArgs().VisitAll(func(key, value []byte) {
		req.Query[string(key)] = string(value)
	})

	ctx.Request.Header.VisitAll(func(key, value []byte) {
		req.Headers[string(key)] = string(value)
	})

	// Let Lua handle everything
	err := worker.HandleRequest(req, resp)
	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
		ctx.SetBodyString(fmt.Sprintf("Internal Server Error: %v", err))
		return
	}

	// Apply response
	ctx.SetStatusCode(resp.StatusCode)
	for key, value := range resp.Headers {
		ctx.Response.Header.Set(key, value)
	}
	if resp.Body != "" {
		ctx.SetBodyString(resp.Body)
	}
}

// RegisterStaticHandler adds a static file handler
func RegisterStaticHandler(urlPrefix, rootPath string) {
	staticMu.Lock()
	defer staticMu.Unlock()

	fs := &fasthttp.FS{
		Root:               rootPath,
		IndexNames:         []string{"index.html"},
		GenerateIndexPages: false,
		Compress:           true,
		AcceptByteRange:    true,
	}

	staticHandlers[urlPrefix] = fs
}