synctv/server/handlers/proxy/cache.go

package proxy

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	json "github.com/json-iterator/go"
	"github.com/zijiren233/gencontainer/dllist"
	"github.com/zijiren233/ksync"
)

// Cache defines the interface for cache implementations
type Cache interface {
	Get(key string) (*CacheItem, bool, error)
	GetAnyWithPrefix(prefix string) (*CacheItem, bool, error)
	Set(key string, data *CacheItem) error
}

// CacheMetadata stores metadata about a cached response
type CacheMetadata struct {
	Headers            http.Header `json:"h,omitempty"`
	ContentType        string      `json:"ct,omitempty"`
	ContentTotalLength int64       `json:"ctl,omitempty"`
}

func (m *CacheMetadata) MarshalBinary() ([]byte, error) {
	return json.Marshal(m)
}

// CacheItem represents a cached response with metadata and data
type CacheItem struct {
	Metadata *CacheMetadata
	Data     []byte
}

// WriteTo implements io.WriterTo to serialize the cache item
func (i *CacheItem) WriteTo(w io.Writer) (int64, error) {
	if w == nil {
		return 0, errors.New("cannot write to nil io.Writer")
	}

	if i.Metadata == nil {
		return 0, errors.New("CacheItem contains nil Metadata")
	}

	metadata, err := i.Metadata.MarshalBinary()
	if err != nil {
		return 0, fmt.Errorf("failed to marshal metadata: %w", err)
	}

	var written int64

	// Write metadata length and content
	if err := binary.Write(w, binary.BigEndian, int64(len(metadata))); err != nil {
		return written, fmt.Errorf("failed to write metadata length: %w", err)
	}
	written += 8

	n, err := w.Write(metadata)
	written += int64(n)
	if err != nil {
		return written, fmt.Errorf("failed to write metadata bytes: %w", err)
	}

	// Write data length and content
	if err := binary.Write(w, binary.BigEndian, int64(len(i.Data))); err != nil {
		return written, fmt.Errorf("failed to write data length: %w", err)
	}
	written += 8

	n, err = w.Write(i.Data)
	written += int64(n)
	if err != nil {
		return written, fmt.Errorf("failed to write data bytes: %w", err)
	}

	return written, nil
}

// ReadFrom implements io.ReaderFrom to deserialize the cache item
func (i *CacheItem) ReadFrom(r io.Reader) (int64, error) {
	if r == nil {
		return 0, errors.New("cannot read from nil io.Reader")
	}

	var read int64

	// Read metadata length and content
	var metadataLen int64
	if err := binary.Read(r, binary.BigEndian, &metadataLen); err != nil {
		return read, fmt.Errorf("failed to read metadata length: %w", err)
	}
	read += 8

	if metadataLen <= 0 {
		return read, fmt.Errorf("metadata length must be positive, got: %d", metadataLen)
	}

	metadata := make([]byte, metadataLen)
	n, err := io.ReadFull(r, metadata)
	read += int64(n)
	if err != nil {
		return read, fmt.Errorf("failed to read metadata bytes: %w", err)
	}

	i.Metadata = new(CacheMetadata)
	if err := json.Unmarshal(metadata, i.Metadata); err != nil {
		return read, fmt.Errorf("failed to unmarshal metadata: %w", err)
	}

	// Read data length and content
	var dataLen int64
	if err := binary.Read(r, binary.BigEndian, &dataLen); err != nil {
		return read, fmt.Errorf("failed to read data length: %w", err)
	}
	read += 8

	if dataLen < 0 {
		return read, fmt.Errorf("data length cannot be negative, got: %d", dataLen)
	}

	i.Data = make([]byte, dataLen)
	n, err = io.ReadFull(r, i.Data)
	read += int64(n)
	if err != nil {
		return read, fmt.Errorf("failed to read data bytes: %w", err)
	}

	return read, nil
}

// TrieNode represents a node in the prefix tree
type TrieNode struct {
	children map[rune]*TrieNode
	key      string
	isEnd    bool
}

func NewTrieNode() *TrieNode {
	return &TrieNode{
		children: make(map[rune]*TrieNode),
	}
}

// MemoryCache implements an in-memory Cache with LRU eviction
type MemoryCache struct {
	m            map[string]*dllist.Element[*cacheEntry]
	lruList      *dllist.Dllist[*cacheEntry]
	prefixTrie   *TrieNode
	capacity     int
	maxSizeBytes int64
	currentSize  int64
	mu           sync.RWMutex
}

type MemoryCacheOption func(*MemoryCache)

func WithMaxSizeBytes(size int64) MemoryCacheOption {
	return func(c *MemoryCache) {
		c.maxSizeBytes = size
	}
}

type cacheEntry struct {
	item *CacheItem
	key  string
	size int64
}

func NewMemoryCache(capacity int, opts ...MemoryCacheOption) *MemoryCache {
	mc := &MemoryCache{
		m:          make(map[string]*dllist.Element[*cacheEntry]),
		lruList:    dllist.New[*cacheEntry](),
		capacity:   capacity,
		prefixTrie: NewTrieNode(),
	}
	for _, opt := range opts {
		opt(mc)
	}
	return mc
}

func (c *MemoryCache) Get(key string) (*CacheItem, bool, error) {
	if key == "" {
		return nil, false, errors.New("cache key cannot be empty")
	}

	c.mu.RLock()
	element, exists := c.m[key]
	if !exists {
		c.mu.RUnlock()
		return nil, false, nil
	}

	// Upgrade to write lock for moving element
	c.mu.RUnlock()
	c.mu.Lock()
	c.lruList.MoveToFront(element)
	item := element.Value.item
	c.mu.Unlock()

	return item, true, nil
}

func (c *MemoryCache) GetAnyWithPrefix(prefix string) (*CacheItem, bool, error) {
	if prefix == "" {
		return nil, false, errors.New("prefix cannot be empty")
	}

	c.mu.RLock()
	defer c.mu.RUnlock()

	// Find matching key in prefix tree
	node := c.prefixTrie
	for _, ch := range prefix {
		if next, ok := node.children[ch]; ok {
			node = next
		} else {
			return nil, false, nil
		}
	}

	// DFS to find first complete key
	var findKey func(*TrieNode) string
	findKey = func(n *TrieNode) string {
		if n.isEnd {
			return n.key
		}
		for _, child := range n.children {
			if key := findKey(child); key != "" {
				return key
			}
		}
		return ""
	}

	if key := findKey(node); key != "" {
		if element, ok := c.m[key]; ok {
			return element.Value.item, true, nil
		}
	}

	return nil, false, nil
}

func (c *MemoryCache) Set(key string, data *CacheItem) error {
	if key == "" {
		return errors.New("cache key cannot be empty")
	}
	if data == nil {
		return errors.New("cannot cache nil CacheItem")
	}

	// Calculate size of new item
	newSize := int64(len(data.Data))
	if data.Metadata != nil {
		metadataBytes, err := data.Metadata.MarshalBinary()
		if err == nil {
			newSize += int64(len(metadataBytes))
		}
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	// Update existing entry if present
	if element, ok := c.m[key]; ok {
		c.currentSize -= element.Value.size
		c.currentSize += newSize
		c.lruList.MoveToFront(element)
		element.Value.item = data
		element.Value.size = newSize
		return nil
	}

	// Evict entries if needed
	for c.lruList.Len() > 0 &&
		((c.capacity > 0 && c.lruList.Len() >= c.capacity) ||
			(c.maxSizeBytes > 0 && c.currentSize+newSize > c.maxSizeBytes)) {
		if back := c.lruList.Back(); back != nil {
			entry := back.Value
			c.currentSize -= entry.size
			delete(c.m, entry.key)
			c.lruList.Remove(back)

			// Remove from prefix tree
			node := c.prefixTrie
			for _, ch := range entry.key {
				node = node.children[ch]
			}
			node.isEnd = false
			node.key = ""
		}
	}

	// Add new entry
	newEntry := &cacheEntry{key: key, item: data, size: newSize}
	element := c.lruList.PushFront(newEntry)
	c.m[key] = element
	c.currentSize += newSize

	// Add to prefix tree
	node := c.prefixTrie
	for _, ch := range key {
		if next, ok := node.children[ch]; ok {
			node = next
		} else {
			node.children[ch] = NewTrieNode()
			node = node.children[ch]
		}
	}
	node.isEnd = true
	node.key = key

	return nil
}

type FileCache struct {
	mu           *ksync.Krwmutex
	memCache     *MemoryCache
	filePath     string
	maxSizeBytes int64
	currentSize  atomic.Int64
	lastCleanup  atomic.Int64
	maxAge       time.Duration
	cleanMu      sync.Mutex
}

type FileCacheOption func(*FileCache)

func WithFileCacheMaxSizeBytes(size int64) FileCacheOption {
	return func(c *FileCache) {
		c.maxSizeBytes = size
	}
}

func WithFileCacheMaxAge(age time.Duration) FileCacheOption {
	return func(c *FileCache) {
		if age > 0 {
			c.maxAge = age
		}
	}
}

func NewFileCache(filePath string, opts ...FileCacheOption) *FileCache {
	fc := &FileCache{
		filePath: filePath,
		mu:       ksync.DefaultKrwmutex(),
		maxAge:   24 * time.Hour, // Default 1 day
		// Initialize memory cache with 1000 items capacity
		memCache: NewMemoryCache(1000, WithMaxSizeBytes(100*1024*1024)), // 100MB memory cache
	}

	for _, opt := range opts {
		opt(fc)
	}

	go fc.periodicCleanup()
	return fc
}

func (c *FileCache) periodicCleanup() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		c.cleanup()
	}
}

func (c *FileCache) cleanup() {
	maxSize := c.maxSizeBytes
	if maxSize <= 0 {
		return
	}

	// Avoid frequent cleanups
	now := time.Now().Unix()
	if now-c.lastCleanup.Load() < 300 {
		return
	}

	c.cleanMu.Lock()
	defer c.cleanMu.Unlock()

	// Double check after acquiring lock
	if now-c.lastCleanup.Load() < 300 {
		return
	}

	entries, err := os.ReadDir(c.filePath)
	if err != nil {
		return
	}

	type fileInfo struct {
		modTime time.Time
		path    string
		size    int64
	}

	var files []fileInfo
	var totalSize int64
	cutoffTime := time.Now().Add(-c.maxAge)

	// Collect file information and remove expired files
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}

		subdir := filepath.Join(c.filePath, entry.Name())
		subEntries, err := os.ReadDir(subdir)
		if err != nil {
			continue
		}

		for _, subEntry := range subEntries {
			info, err := subEntry.Info()
			if err != nil {
				continue
			}

			fullPath := filepath.Join(subdir, subEntry.Name())

			// Remove expired files
			if info.ModTime().Before(cutoffTime) {
				os.Remove(fullPath)
				continue
			}

			files = append(files, fileInfo{
				path:    fullPath,
				size:    info.Size(),
				modTime: info.ModTime(),
			})
			totalSize += info.Size()
		}
	}

	// If under size limit, just update size and return
	if totalSize <= maxSize {
		c.currentSize.Store(totalSize)
		c.lastCleanup.Store(now)
		return
	}

	// Sort by modification time (oldest first) and remove until under limit
	sort.Slice(files, func(i, j int) bool {
		return files[i].modTime.Before(files[j].modTime)
	})

	for _, file := range files {
		if totalSize <= maxSize {
			break
		}
		if err := os.Remove(file.path); err == nil {
			totalSize -= file.size
		}
	}

	c.currentSize.Store(totalSize)
	c.lastCleanup.Store(now)
}

func (c *FileCache) Get(key string) (*CacheItem, bool, error) {
	if key == "" {
		return nil, false, errors.New("cache key cannot be empty")
	}

	// Try memory cache first
	if item, found, err := c.memCache.Get(key); err == nil && found {
		return item, true, nil
	}

	prefix := string(key[0])
	filePath := filepath.Join(c.filePath, prefix, key)

	c.mu.RLock(key)
	defer c.mu.RUnlock(key)

	file, err := os.OpenFile(filePath, os.O_RDONLY, 0o644)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, false, nil
		}
		return nil, false, fmt.Errorf("failed to open cache file: %w", err)
	}
	defer file.Close()

	// Check if file is expired
	if info, err := file.Stat(); err == nil {
		if time.Since(info.ModTime()) > c.maxAge {
			os.Remove(filePath)
			return nil, false, nil
		}
	}

	item := &CacheItem{}
	if _, err := item.ReadFrom(file); err != nil {
		return nil, false, fmt.Errorf("failed to read cache item: %w", err)
	}

	// Store in memory cache
	err = c.memCache.Set(key, item)
	if err != nil {
		return nil, false, fmt.Errorf("failed to set cache item: %w", err)
	}

	return item, true, nil
}

func (c *FileCache) GetAnyWithPrefix(prefix string) (*CacheItem, bool, error) {
	if prefix == "" {
		return nil, false, errors.New("prefix cannot be empty")
	}

	// Try memory cache first
	if item, found, err := c.memCache.GetAnyWithPrefix(prefix); err == nil && found {
		return item, true, nil
	}

	prefixDir := string(prefix[0])
	dirPath := filepath.Join(c.filePath, prefixDir)

	entries, err := os.ReadDir(dirPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, false, nil
		}
		return nil, false, fmt.Errorf("failed to read directory: %w", err)
	}

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		name := entry.Name()
		if len(name) >= len(prefix) && name[:len(prefix)] == prefix {
			item, found, err := c.Get(name)
			if err == nil && found {
				return item, true, nil
			}
		}
	}

	return nil, false, nil
}

func (c *FileCache) Set(key string, data *CacheItem) error {
	if key == "" {
		return errors.New("cache key cannot be empty")
	}
	if data == nil {
		return errors.New("cannot cache nil CacheItem")
	}

	// Store in memory cache first
	if err := c.memCache.Set(key, data); err != nil {
		return err
	}

	// Check and cleanup if needed
	maxSize := c.maxSizeBytes
	if maxSize > 0 {
		newSize := int64(len(data.Data))
		if data.Metadata != nil {
			if metadataBytes, err := data.Metadata.MarshalBinary(); err == nil {
				newSize += int64(len(metadataBytes))
			}
		}
		if c.currentSize.Load()+newSize > maxSize {
			c.cleanup()
		}
	}

	prefix := string(key[0])
	dirPath := filepath.Join(c.filePath, prefix)
	filePath := filepath.Join(dirPath, key)

	c.mu.Lock(key)
	defer c.mu.Unlock(key)

	if err := os.MkdirAll(dirPath, 0o755); err != nil {
		return fmt.Errorf("failed to create cache directory: %w", err)
	}

	file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
	if err != nil {
		return fmt.Errorf("failed to create cache file: %w", err)
	}
	defer file.Close()

	if _, err := data.WriteTo(file); err != nil {
		return fmt.Errorf("failed to write cache item: %w", err)
	}

	if info, err := file.Stat(); err == nil {
		c.currentSize.Add(info.Size())
	}

	return nil
}