Moonshark/modules/http/http.go

package http

import (
	"Moonshark/metadata"
	"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{
		Name:         "Moonshark/" + metadata.Version,
		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 {
	StopAllServers()
	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)
	noCache := s.ToBoolean(3)

	// 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, noCache)
	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())

	// Fast path for likely static files (has extension)
	if isLikelyStaticFile(path) && tryStaticHandler(ctx, path) {
		return
	}

	// Try Lua routing
	globalMu.RLock()
	pool := globalWorkerPool
	globalMu.RUnlock()

	if pool != nil {
		worker := pool.Get()
		if worker != nil {
			defer pool.Put(worker)

			req := GetRequest()
			defer PutRequest(req)

			resp := GetResponse()
			defer PutResponse(resp)

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

			for key, value := range ctx.QueryArgs().All() {
				req.Query[string(key)] = string(value)
			}

			for key, value := range ctx.Request.Header.All() {
				req.Headers[string(key)] = string(value)
			}

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

			// If Lua found a route, use it
			if resp.StatusCode != 404 {
				ctx.SetStatusCode(resp.StatusCode)
				for key, value := range resp.Headers {
					ctx.Response.Header.Set(key, value)
				}
				if resp.Body != "" {
					ctx.SetBodyString(resp.Body)
				}
				return
			}
		}
	}

	if !isLikelyStaticFile(path) && tryStaticHandler(ctx, path) {
		return
	}

	ctx.SetStatusCode(fasthttp.StatusNotFound)
	ctx.SetBodyString("Not Found")
}

func tryStaticHandler(ctx *fasthttp.RequestCtx, path string) bool {
	staticMu.RLock()
	defer staticMu.RUnlock()

	for prefix, fs := range staticHandlers {
		if after, ok := strings.CutPrefix(path, prefix); ok {
			ctx.Request.URI().SetPath(after)
			fs.NewRequestHandler()(ctx)
			return true
		}
	}
	return false
}

func isLikelyStaticFile(path string) bool {
	// Check for file extension
	lastSlash := strings.LastIndex(path, "/")
	lastDot := strings.LastIndex(path, ".")
	return lastDot > lastSlash && lastDot != -1
}

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

	var cacheDuration time.Duration
	var compress bool
	if noCache {
		cacheDuration = 0
		compress = false
	} else {
		cacheDuration = 3600 * time.Second
		compress = true
	}

	fs := &fasthttp.FS{
		Root:               rootPath,
		CompressRoot:       rootPath + "/.cache",
		IndexNames:         []string{"index.html"},
		GenerateIndexPages: false,
		Compress:           compress,
		CompressBrotli:     compress,
		CompressZstd:       compress,
		CacheDuration:      cacheDuration,
		AcceptByteRange:    true,
		PathNotFound: func(ctx *fasthttp.RequestCtx) {
			path := ctx.Path()
			fmt.Printf("404 not found: %s\n", path)
			ctx.SetStatusCode(fasthttp.StatusNotFound)
			ctx.SetBodyString("404 not found")
		},
	}

	staticHandlers[urlPrefix] = fs
}

func StopAllServers() {
	globalMu.Lock()
	defer globalMu.Unlock()

	// Start shutting down both in parallel
	var wg sync.WaitGroup

	// Close worker pool in parallel
	if globalWorkerPool != nil {
		wg.Add(1)
		pool := globalWorkerPool
		globalWorkerPool = nil
		go func() {
			defer wg.Done()
			pool.Close()
		}()
	}

	// Shutdown server with 100ms timeout
	if globalServer != nil {
		wg.Add(1)
		server := globalServer
		globalServer = nil
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
			defer cancel()
			if err := server.ShutdownWithContext(ctx); err != nil {
				// Force close if graceful shutdown times out
				server.CloseOnShutdown = true
				_ = server.Shutdown()
			}
		}()
	}

	serverRunning = false

	// Wait for both to complete
	wg.Wait()
}