Moonshark/core/routers/StaticRouter.go

package routers

import (
	"bytes"
	"compress/gzip"
	"container/list"
	"errors"
	"mime"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.sharkk.net/Sky/Moonshark/core/logger"
)

// CacheEntry represents a cached file with metadata
type CacheEntry struct {
	Path           string // URL path
	GzippedContent []byte // Gzipped content
	Size           int    // Original size
	GzippedSize    int    // Compressed size
	ModTime        int64  // Modification time
	ContentType    string // Content type
}

// StaticRouter is a caching router for static files
// It provides an LRU cache for gzipped static assets
type StaticRouter struct {
	rootDir       string                   // Root directory containing files
	cache         map[string]*list.Element // Cache map (URL path -> list element)
	lruList       *list.List               // LRU tracking list
	mu            sync.RWMutex             // Lock for concurrent access
	maxItems      int                      // Maximum number of items in cache
	maxItemSize   int                      // Maximum size per item (gzipped)
	totalCapacity int                      // Total cache capacity in bytes
	currentSize   int                      // Current cache size in bytes
	fileServer    http.Handler             // Underlying file server
	bufferPool    sync.Pool                // Buffer pool for compression
	urlPrefix     string                   // URL prefix for static assets
	log           bool                     // Whether to log debug info
}

// NewStaticRouterWithLogger creates a new StaticRouter instance
func NewStaticRouter(rootDir string) (*StaticRouter, error) {
	// Verify root directory exists
	info, err := os.Stat(rootDir)
	if err != nil {
		return nil, err
	}
	if !info.IsDir() {
		return nil, errors.New("root path is not a directory")
	}

	// Create the router with default settings
	r := &StaticRouter{
		rootDir:       rootDir,
		cache:         make(map[string]*list.Element),
		lruList:       list.New(),
		maxItems:      100,      // Default: cache 100 files
		maxItemSize:   1 << 20,  // Default: 1MB per file (gzipped)
		totalCapacity: 20 << 20, // Default: 20MB total cache
		fileServer:    http.FileServer(http.Dir(rootDir)),
		urlPrefix:     "/static", // Default prefix for static assets
		log:           false,     // Debug logging off by default
		bufferPool: sync.Pool{
			New: func() any {
				return new(bytes.Buffer)
			},
		},
	}

	// Initialize mime package with common types
	mime.AddExtensionType(".js", "application/javascript")
	mime.AddExtensionType(".css", "text/css")
	mime.AddExtensionType(".svg", "image/svg+xml")

	return r, nil
}

// SetMaxItems sets the maximum number of items in the cache
func (r *StaticRouter) SetMaxItems(n int) {
	if n <= 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.maxItems = n
}

// SetMaxItemSize sets the maximum size per cached item (in bytes)
func (r *StaticRouter) SetMaxItemSize(n int) {
	if n <= 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.maxItemSize = n
}

// SetTotalCapacity sets the total cache capacity (in bytes)
func (r *StaticRouter) SetTotalCapacity(n int) {
	if n <= 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.totalCapacity = n
}

// EnableDebugLog enables debug logging
func (r *StaticRouter) EnableDebugLog() {
	r.log = true
}

// DisableDebugLog disables debug logging
func (r *StaticRouter) DisableDebugLog() {
	r.log = false
}

// SetURLPrefix sets the URL prefix for static assets
func (r *StaticRouter) SetURLPrefix(prefix string) {
	if !strings.HasPrefix(prefix, "/") {
		prefix = "/" + prefix
	}
	r.urlPrefix = prefix
}

// ServeHTTP implements http.Handler interface
func (r *StaticRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Check if path starts with the prefix
	if !strings.HasPrefix(req.URL.Path, r.urlPrefix) {
		http.NotFound(w, req)
		return
	}

	// Strip prefix to get the file path
	origPath := req.URL.Path
	fileURLPath := strings.TrimPrefix(req.URL.Path, r.urlPrefix)

	// Make sure path starts with /
	if !strings.HasPrefix(fileURLPath, "/") {
		fileURLPath = "/" + fileURLPath
	}

	// Check for directory access without trailing slash
	if strings.HasSuffix(fileURLPath, "/index.html") {
		dirPath := fileURLPath[:len(fileURLPath)-10] // remove "/index.html"
		filePath := filepath.Join(r.rootDir, filepath.FromSlash(strings.TrimPrefix(dirPath, "/")))
		if info, err := os.Stat(filePath); err == nil && info.IsDir() {
			// This is a directory with index.html, handle it directly
			fileURLPath = dirPath + "/"
		}
	}

	// Check if client accepts gzip encoding
	acceptsGzip := strings.Contains(req.Header.Get("Accept-Encoding"), "gzip")

	// Copy the original request for modification
	newReq := *req
	newReq.URL = new(url.URL)
	*newReq.URL = *req.URL
	newReq.URL.Path = fileURLPath

	// Try to serve from cache if client accepts gzip
	if acceptsGzip && r.serveFromCache(w, &newReq, fileURLPath) {
		if r.log {
			logger.Debug("[StaticRouter] CACHE HIT: %s", origPath)
		}
		return
	}

	if r.log {
		logger.Debug("[StaticRouter] CACHE MISS: %s", origPath)
	}

	// Fall back to standard file serving
	wrappedWriter := httptest.NewRecorder()
	r.fileServer.ServeHTTP(wrappedWriter, &newReq)

	// Check if we got a redirect - might need to add the prefix back
	if wrappedWriter.Code == http.StatusMovedPermanently || wrappedWriter.Code == http.StatusPermanentRedirect {
		location := wrappedWriter.Header().Get("Location")
		if location != "" && !strings.HasPrefix(location, r.urlPrefix) {
			// Prepend our prefix to the redirect location
			newLocation := r.urlPrefix
			if !strings.HasPrefix(location, "/") {
				newLocation += "/"
			}
			newLocation += strings.TrimPrefix(location, "/")
			wrappedWriter.Header().Set("Location", newLocation)
		}
	}

	// Copy the response from the recorder to the real response writer
	for k, v := range wrappedWriter.Header() {
		w.Header()[k] = v
	}
	w.WriteHeader(wrappedWriter.Code)
	w.Write(wrappedWriter.Body.Bytes())

	// Try to cache the file for next time if client accepts gzip
	if acceptsGzip && wrappedWriter.Code == http.StatusOK {
		filePath := filepath.Join(r.rootDir, filepath.FromSlash(strings.TrimPrefix(fileURLPath, "/")))
		// Cache synchronously for tests
		r.cacheFile(fileURLPath, filePath)
	}
}

// Match finds a file path for the given URL path
func (r *StaticRouter) Match(urlPath string) (string, bool) {
	// Check if path starts with the prefix
	if !strings.HasPrefix(urlPath, r.urlPrefix) {
		return "", false
	}

	// Strip prefix
	urlPath = strings.TrimPrefix(urlPath, r.urlPrefix)

	// Make sure path starts with /
	if !strings.HasPrefix(urlPath, "/") {
		urlPath = "/" + urlPath
	}

	filePath := filepath.Join(r.rootDir, filepath.FromSlash(strings.TrimPrefix(urlPath, "/")))
	_, err := os.Stat(filePath)

	if r.log && err == nil {
		logger.Debug("[StaticRouter] MATCH: %s -> %s", urlPath, filePath)
	}

	return filePath, err == nil
}

// serveFromCache tries to serve a file from cache
// Returns true if successful, false otherwise
func (r *StaticRouter) serveFromCache(w http.ResponseWriter, req *http.Request, urlPath string) bool {
	// Check cache first with read lock
	r.mu.RLock()
	elem, found := r.cache[urlPath]
	if !found {
		r.mu.RUnlock()
		return false
	}

	// Get cache entry
	entry := elem.Value.(*CacheEntry)
	content := entry.GzippedContent
	contentType := entry.ContentType
	modTime := time.Unix(entry.ModTime, 0)
	r.mu.RUnlock()

	// Update LRU order with write lock
	r.mu.Lock()
	r.lruList.MoveToFront(elem)
	r.mu.Unlock()

	// Check if client cache is still valid (If-Modified-Since)
	if !isModified(req, modTime) {
		w.Header().Set("Content-Type", contentType)
		w.Header().Set("Last-Modified", modTime.UTC().Format(http.TimeFormat))
		w.Header().Set("Vary", "Accept-Encoding")
		w.WriteHeader(http.StatusNotModified)
		return true
	}

	// Set appropriate headers
	w.Header().Set("Content-Encoding", "gzip")
	w.Header().Set("Content-Type", contentType)
	w.Header().Set("Content-Length", strconv.Itoa(len(content)))
	w.Header().Set("Vary", "Accept-Encoding")
	w.Header().Set("Last-Modified", modTime.UTC().Format(http.TimeFormat))

	// Write the content directly instead of using ServeContent
	w.WriteHeader(http.StatusOK)
	w.Write(content)
	return true
}

// isModified checks if the file has been modified since the client's last request
func isModified(req *http.Request, modTime time.Time) bool {
	// Parse If-Modified-Since header
	if ims := req.Header.Get("If-Modified-Since"); ims != "" {
		t, err := http.ParseTime(ims)
		if err == nil && !modTime.After(t.Add(time.Second)) {
			return false
		}
	}
	return true
}

// cacheFile adds a file to the cache
func (r *StaticRouter) cacheFile(urlPath, filePath string) {
	// Stat the file
	fileInfo, err := os.Stat(filePath)
	if err != nil || fileInfo.IsDir() {
		return
	}

	// Don't cache files that are too large
	if fileInfo.Size() > int64(r.maxItemSize*2) {
		return
	}

	// Read file
	data, err := os.ReadFile(filePath)
	if err != nil {
		return
	}

	// Compress data using buffer from pool
	buf := r.bufferPool.Get().(*bytes.Buffer)
	buf.Reset()
	defer r.bufferPool.Put(buf)

	gzWriter, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	if err != nil {
		return
	}

	if _, err := gzWriter.Write(data); err != nil {
		gzWriter.Close()
		return
	}
	if err := gzWriter.Close(); err != nil {
		return
	}

	gzippedData := make([]byte, buf.Len())
	copy(gzippedData, buf.Bytes())

	// Don't cache if compressed size is too large
	if len(gzippedData) > r.maxItemSize {
		return
	}

	// Get content type by extension or detection
	contentType := getMimeType(filePath, data)

	// Update cache
	r.mu.Lock()
	defer r.mu.Unlock()

	// Check if already in cache
	if elem, exists := r.cache[urlPath]; exists {
		// Update existing entry
		oldEntry := elem.Value.(*CacheEntry)
		r.currentSize -= oldEntry.GzippedSize

		newEntry := &CacheEntry{
			Path:           urlPath,
			GzippedContent: gzippedData,
			Size:           len(data),
			GzippedSize:    len(gzippedData),
			ModTime:        fileInfo.ModTime().Unix(),
			ContentType:    contentType,
		}

		elem.Value = newEntry
		r.lruList.MoveToFront(elem)
		r.currentSize += newEntry.GzippedSize
		return
	}

	// Make room in cache if needed
	for r.lruList.Len() >= r.maxItems ||
		(r.currentSize+len(gzippedData) > r.totalCapacity && r.lruList.Len() > 0) {

		// Remove least recently used item
		elem := r.lruList.Back()
		if elem == nil {
			break
		}

		entry := elem.Value.(*CacheEntry)
		r.lruList.Remove(elem)
		delete(r.cache, entry.Path)
		r.currentSize -= entry.GzippedSize
	}

	// Add new item to cache
	entry := &CacheEntry{
		Path:           urlPath,
		GzippedContent: gzippedData,
		Size:           len(data),
		GzippedSize:    len(gzippedData),
		ModTime:        fileInfo.ModTime().Unix(),
		ContentType:    contentType,
	}

	elem := r.lruList.PushFront(entry)
	r.cache[urlPath] = elem
	r.currentSize += entry.GzippedSize
}

// getMimeType returns the content type for a file
func getMimeType(filePath string, data []byte) string {
	// Try to get content type from extension first
	ext := strings.ToLower(filepath.Ext(filePath))
	if mimeType := mime.TypeByExtension(ext); mimeType != "" {
		return mimeType
	}

	// Fall back to detection
	return http.DetectContentType(data)
}

// Refresh clears the cache
func (r *StaticRouter) Refresh() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.cache = make(map[string]*list.Element)
	r.lruList.Init()
	r.currentSize = 0
	return nil
}

// PreloadCommonFiles loads common static file types into the cache
func (r *StaticRouter) PreloadCommonFiles() {
	// Common file extensions to preload
	extensions := map[string]bool{
		".css":  true,
		".js":   true,
		".svg":  true,
		".ico":  true,
		".png":  true,
		".jpg":  true,
		".jpeg": true,
		".gif":  true,
	}

	if r.log {
		logger.Debug("[StaticRouter] Preloading common files from %s", r.rootDir)
	}

	count := 0

	// Walk the directory
	_ = filepath.Walk(r.rootDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}

		// Check extension and file size
		ext := strings.ToLower(filepath.Ext(path))
		if !extensions[ext] || info.Size() > int64(r.maxItemSize*2) {
			return nil
		}

		// Get URL path
		relPath, err := filepath.Rel(r.rootDir, path)
		if err != nil {
			return nil
		}

		// Don't include prefix here - will be applied in ServeHTTP
		urlPath := "/" + strings.ReplaceAll(relPath, "\\", "/")

		// Cache the file synchronously for tests
		r.cacheFile(urlPath, path)
		count++

		return nil
	})

	if r.log {
		logger.Debug("[StaticRouter] Preloaded %d files", count)
	}
}

// GetStats returns cache statistics
func (r *StaticRouter) GetStats() map[string]any {
	r.mu.RLock()
	defer r.mu.RUnlock()

	return map[string]any{
		"items":         r.lruList.Len(),
		"maxItems":      r.maxItems,
		"currentSize":   r.currentSize,
		"totalCapacity": r.totalCapacity,
		"usagePercent":  float64(r.currentSize) * 100 / float64(r.totalCapacity),
	}
}