torrent-gateway/internal/proxy/smart_proxy.go

package proxy

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"git.sovbit.dev/enki/torrentGateway/internal/config"
	"git.sovbit.dev/enki/torrentGateway/internal/storage"
)

// SmartProxy provides intelligent proxy functionality for serving chunked files
type SmartProxy struct {
	storage      *storage.Backend
	gatewayURL   string
	cache        *LRUCache
	config       *config.Config
	mu           sync.RWMutex
}

// NewSmartProxy creates a new smart proxy instance
func NewSmartProxy(storage *storage.Backend, cfg *config.Config) *SmartProxy {
	gatewayURL := fmt.Sprintf("http://localhost:%d", cfg.Gateway.Port)
	cache := NewLRUCache(cfg.Proxy.CacheSize, cfg.Proxy.CacheMaxAge)
	
	return &SmartProxy{
		storage:    storage,
		gatewayURL: gatewayURL,
		cache:      cache,
		config:     cfg,
	}
}

// ServeBlob attempts to serve a blob by hash, reassembling from chunks if necessary
func (p *SmartProxy) ServeBlob(w http.ResponseWriter, hash string) error {
	// First check cache
	if cachedData := p.cache.Get(hash); cachedData != nil {
		log.Printf("Serving cached reassembled file for hash: %s", hash)
		p.serveCachedData(w, hash, cachedData)
		return nil
	}
	
	// Check if this hash exists as chunked file in metadata
	metadata, err := p.storage.GetFileMetadata(hash)
	if err != nil {
		return fmt.Errorf("error checking metadata for hash %s: %v", hash, err)
	}
	
	if metadata == nil {
		return fmt.Errorf("hash %s not found as chunked file", hash)
	}
	
	// Only proceed if this is a torrent/chunked file
	if metadata.StorageType != "torrent" {
		return fmt.Errorf("hash %s is not a chunked file (storage type: %s)", hash, metadata.StorageType)
	}
	
	// Get chunk hashes for this file
	chunkHashes, err := p.storage.GetChunkHashes(hash)
	if err != nil {
		return fmt.Errorf("error getting chunk hashes for %s: %v", hash, err)
	}
	
	if len(chunkHashes) == 0 {
		return fmt.Errorf("no chunks found for file %s", hash)
	}
	
	log.Printf("Found chunked file for hash %s, reassembling %d chunks", hash, len(chunkHashes))
	
	// Reassemble the file from chunks
	reassembledData, err := p.reassembleFile(metadata, chunkHashes)
	if err != nil {
		return fmt.Errorf("error reassembling file %s: %v", hash, err)
	}
	
	// Cache the reassembled data
	p.cache.Put(hash, &CachedBlob{
		Data:     reassembledData,
		MimeType: metadata.ContentType,
		Size:     metadata.Size,
		Hash:     hash,
	})
	
	// Serve the reassembled data
	w.Header().Set("Content-Type", metadata.ContentType)
	w.Header().Set("Content-Length", fmt.Sprintf("%d", len(reassembledData)))
	w.Header().Set("ETag", fmt.Sprintf(`"%s"`, hash))
	w.Header().Set("Cache-Control", "public, max-age=31536000")
	
	if _, err := w.Write(reassembledData); err != nil {
		return fmt.Errorf("error writing response: %v", err)
	}
	
	log.Printf("Successfully served reassembled file for hash: %s (%d bytes)", hash, len(reassembledData))
	return nil
}

// reassembleFile reassembles a file from its chunks
func (p *SmartProxy) reassembleFile(metadata *storage.FileMetadata, chunkHashes []string) ([]byte, error) {
	if len(chunkHashes) == 0 {
		return nil, fmt.Errorf("no chunks found in metadata")
	}
	
	var buf bytes.Buffer
	buf.Grow(int(metadata.Size)) // Pre-allocate buffer
	
	// Process chunks in order
	for i, chunkHash := range chunkHashes {
		chunkData, err := p.storage.GetChunkData(chunkHash)
		if err != nil {
			return nil, fmt.Errorf("error getting chunk %d (%s): %v", i, chunkHash, err)
		}
		
		if chunkData == nil {
			return nil, fmt.Errorf("chunk %d (%s) not found", i, chunkHash)
		}
		
		if _, err := buf.Write(chunkData); err != nil {
			return nil, fmt.Errorf("error writing chunk %d to buffer: %v", i, err)
		}
	}
	
	return buf.Bytes(), nil
}

// serveCachedData serves cached blob data
func (p *SmartProxy) serveCachedData(w http.ResponseWriter, hash string, cached *CachedBlob) {
	w.Header().Set("Content-Type", cached.MimeType)
	w.Header().Set("Content-Length", fmt.Sprintf("%d", len(cached.Data)))
	w.Header().Set("ETag", fmt.Sprintf(`"%s"`, hash))
	w.Header().Set("Cache-Control", "public, max-age=31536000")
	w.Header().Set("X-Proxy-Cache", "HIT")
	
	w.Write(cached.Data)
}

// CachedBlob represents a cached reassembled blob
type CachedBlob struct {
	Data     []byte
	MimeType string
	Size     int64
	Hash     string
	CachedAt time.Time
}

// LRUCache implements a simple LRU cache for reassembled blobs
type LRUCache struct {
	capacity   int
	maxAge     time.Duration
	cache      map[string]*CacheEntry
	lruList    []*CacheEntry
	mu         sync.RWMutex
}

// CacheEntry represents an entry in the cache
type CacheEntry struct {
	Key      string
	Value    *CachedBlob
	CachedAt time.Time
}

// NewLRUCache creates a new LRU cache
func NewLRUCache(capacity int, maxAge time.Duration) *LRUCache {
	if capacity <= 0 {
		capacity = 100 // Default capacity
	}
	if maxAge <= 0 {
		maxAge = 1 * time.Hour // Default max age
	}
	
	return &LRUCache{
		capacity: capacity,
		maxAge:   maxAge,
		cache:    make(map[string]*CacheEntry),
		lruList:  make([]*CacheEntry, 0, capacity),
	}
}

// Get retrieves a value from the cache
func (c *LRUCache) Get(key string) *CachedBlob {
	c.mu.Lock()
	defer c.mu.Unlock()
	
	entry, exists := c.cache[key]
	if !exists {
		return nil
	}
	
	// Check if entry is expired
	if time.Since(entry.CachedAt) > c.maxAge {
		c.removeEntry(key)
		return nil
	}
	
	// Move to front (most recently used)
	c.moveToFront(entry)
	return entry.Value
}

// Put adds a value to the cache
func (c *LRUCache) Put(key string, value *CachedBlob) {
	c.mu.Lock()
	defer c.mu.Unlock()
	
	// Check if entry already exists
	if entry, exists := c.cache[key]; exists {
		entry.Value = value
		entry.CachedAt = time.Now()
		c.moveToFront(entry)
		return
	}
	
	// Create new entry
	entry := &CacheEntry{
		Key:      key,
		Value:    value,
		CachedAt: time.Now(),
	}
	
	// Check capacity
	if len(c.cache) >= c.capacity {
		c.evictLRU()
	}
	
	// Add to cache
	c.cache[key] = entry
	c.lruList = append([]*CacheEntry{entry}, c.lruList...)
}

// moveToFront moves an entry to the front of the LRU list
func (c *LRUCache) moveToFront(entry *CacheEntry) {
	// Find and remove entry from current position
	for i, e := range c.lruList {
		if e == entry {
			c.lruList = append(c.lruList[:i], c.lruList[i+1:]...)
			break
		}
	}
	
	// Add to front
	c.lruList = append([]*CacheEntry{entry}, c.lruList...)
}

// evictLRU removes the least recently used entry
func (c *LRUCache) evictLRU() {
	if len(c.lruList) == 0 {
		return
	}
	
	// Remove last entry (LRU)
	lru := c.lruList[len(c.lruList)-1]
	c.lruList = c.lruList[:len(c.lruList)-1]
	delete(c.cache, lru.Key)
	
	log.Printf("Evicted cached blob: %s", lru.Key)
}

// removeEntry removes an entry from the cache
func (c *LRUCache) removeEntry(key string) {
	entry, exists := c.cache[key]
	if !exists {
		return
	}
	
	// Remove from cache map
	delete(c.cache, key)
	
	// Remove from LRU list
	for i, e := range c.lruList {
		if e == entry {
			c.lruList = append(c.lruList[:i], c.lruList[i+1:]...)
			break
		}
	}
}

// CleanExpired removes expired entries from the cache
func (c *LRUCache) CleanExpired() {
	c.mu.Lock()
	defer c.mu.Unlock()
	
	now := time.Now()
	var toRemove []string
	
	for key, entry := range c.cache {
		if now.Sub(entry.CachedAt) > c.maxAge {
			toRemove = append(toRemove, key)
		}
	}
	
	for _, key := range toRemove {
		c.removeEntry(key)
	}
	
	if len(toRemove) > 0 {
		log.Printf("Cleaned %d expired cache entries", len(toRemove))
	}
}

// Size returns the current cache size
func (c *LRUCache) Size() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.cache)
}