// vfs/memory/memory.go
package memory

import (
	"bytes"
	"fmt"
	"io"
	"s1d3sw1ped/steamcache2/vfs"
	"s1d3sw1ped/steamcache2/vfs/locks"
	"s1d3sw1ped/steamcache2/vfs/lru"
	"s1d3sw1ped/steamcache2/vfs/types"
	"s1d3sw1ped/steamcache2/vfs/vfserror"
	"sort"
	"strings"
	"sync"
	"time"
)

// maxEvictBatch bounds the candidate snapshot during RLock/Lock collect in Evict*.
// Prevents holding lock for unbounded time under extreme pressure.
const maxEvictBatch = 4096

// Ensure MemoryFS implements VFS.
var _ vfs.VFS = (*MemoryFS)(nil)

// MemoryFS is an in-memory virtual file system
type MemoryFS struct {
	data        map[string]*bytes.Buffer
	info        map[string]*types.FileInfo
	capacity    int64
	size        int64
	mu          sync.RWMutex
	keyLocks    []sync.Map // Sharded lock pools for better concurrency
	LRU         *lru.LRUList[*types.FileInfo]
	timeUpdater *types.BatchedTimeUpdate // Batched time updates for better performance
}

// New creates a new MemoryFS
func New(capacity int64) (*MemoryFS, error) {
	if capacity <= 0 {
		return nil, fmt.Errorf("memory capacity must be greater than 0")
	}

	// Initialize sharded locks
	keyLocks := make([]sync.Map, locks.NumLockShards)

	return &MemoryFS{
		data:        make(map[string]*bytes.Buffer),
		info:        make(map[string]*types.FileInfo),
		capacity:    capacity,
		size:        0,
		keyLocks:    keyLocks,
		LRU:         lru.NewLRUList[*types.FileInfo](),
		timeUpdater: types.NewBatchedTimeUpdate(100 * time.Millisecond), // Update time every 100ms
	}, nil
}

// Name returns the name of this VFS
func (m *MemoryFS) Name() string {
	return "MemoryFS"
}

// Size returns the current size
func (m *MemoryFS) Size() int64 {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.size
}

// Capacity returns the maximum capacity
func (m *MemoryFS) Capacity() int64 {
	return m.capacity
}

// GetFragmentationStats returns memory fragmentation statistics
func (m *MemoryFS) GetFragmentationStats() map[string]interface{} {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var totalCapacity int64
	var totalUsed int64
	var bufferCount int

	for _, buffer := range m.data {
		totalCapacity += int64(buffer.Cap())
		totalUsed += int64(buffer.Len())
		bufferCount++
	}

	fragmentationRatio := float64(0)
	if totalCapacity > 0 {
		fragmentationRatio = float64(totalCapacity-totalUsed) / float64(totalCapacity)
	}

	return map[string]interface{}{
		"buffer_count":        bufferCount,
		"total_capacity":      totalCapacity,
		"total_used":          totalUsed,
		"fragmentation_ratio": fragmentationRatio,
		"average_buffer_size": float64(totalUsed) / float64(bufferCount),
	}
}

// getKeyLock returns a lock for the given key using sharding
func (m *MemoryFS) getKeyLock(key string) *sync.RWMutex {
	return locks.GetKeyLock(m.keyLocks, key)
}

// Create creates a new file
func (m *MemoryFS) Create(key string, size int64) (io.WriteCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	m.mu.Lock()
	// Check if file already exists and handle overwrite
	if fi, exists := m.info[key]; exists {
		m.size -= fi.Size
		m.LRU.Remove(key)
		delete(m.info, key)
		delete(m.data, key)
	}

	buffer := &bytes.Buffer{}
	m.data[key] = buffer
	fi := types.NewFileInfo(key, size)
	m.info[key] = fi
	m.LRU.Add(key, fi)
	// Initialize access time with current time
	fi.UpdateAccessBatched(m.timeUpdater)
	m.size += size
	m.mu.Unlock()

	return &memoryWriteCloser{
		buffer: buffer,
		memory: m,
		key:    key,
	}, nil
}

// memoryWriteCloser implements io.WriteCloser for memory files
type memoryWriteCloser struct {
	buffer *bytes.Buffer
	memory *MemoryFS
	key    string
}

func (mwc *memoryWriteCloser) Write(p []byte) (n int, err error) {
	return mwc.buffer.Write(p)
}

func (mwc *memoryWriteCloser) Close() error {
	// Update the actual size in FileInfo
	mwc.memory.mu.Lock()
	if fi, exists := mwc.memory.info[mwc.key]; exists {
		actualSize := int64(mwc.buffer.Len())
		sizeDiff := actualSize - fi.Size
		fi.Size = actualSize
		mwc.memory.size += sizeDiff
	}
	mwc.memory.mu.Unlock()
	return nil
}

// Open opens a file for reading
func (m *MemoryFS) Open(key string) (io.ReadCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.Lock()
	fi, exists := m.info[key]
	if !exists {
		m.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}
	fi.UpdateAccessBatched(m.timeUpdater)
	m.LRU.MoveToFront(key, m.timeUpdater)

	buffer, exists := m.data[key]
	if !exists {
		m.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}

	// Use zero-copy approach - return reader that reads directly from buffer
	m.mu.Unlock()

	return &memoryReadCloser{
		buffer: buffer,
		offset: 0,
	}, nil
}

// memoryReadCloser implements io.ReadCloser for memory files with zero-copy optimization
type memoryReadCloser struct {
	buffer *bytes.Buffer
	offset int64
}

func (mrc *memoryReadCloser) Read(p []byte) (n int, err error) {
	if mrc.offset >= int64(mrc.buffer.Len()) {
		return 0, io.EOF
	}

	// Zero-copy read directly from buffer
	available := mrc.buffer.Len() - int(mrc.offset)
	toRead := len(p)
	if toRead > available {
		toRead = available
	}

	// Read directly from buffer without copying
	data := mrc.buffer.Bytes()
	copy(p, data[mrc.offset:mrc.offset+int64(toRead)])
	mrc.offset += int64(toRead)

	return toRead, nil
}

func (mrc *memoryReadCloser) Close() error {
	return nil
}

// Delete removes a file
func (m *MemoryFS) Delete(key string) error {
	if key == "" {
		return vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	m.mu.Lock()
	fi, exists := m.info[key]
	if !exists {
		m.mu.Unlock()
		return vfserror.ErrNotFound
	}
	m.size -= fi.Size
	m.LRU.Remove(key)
	delete(m.info, key)
	delete(m.data, key)
	m.mu.Unlock()

	return nil
}

// Stat returns file information
func (m *MemoryFS) Stat(key string) (*types.FileInfo, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := m.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	m.mu.RLock()
	defer m.mu.RUnlock()

	if fi, ok := m.info[key]; ok {
		return fi, nil
	}

	return nil, vfserror.ErrNotFound
}

// EvictLRU evicts the least recently used files to free up space
// Collect under short exclusive Lock (to serialize concurrent EvictLRU on the unsynchronized LRUList),
// then batch delete under WLock. Regular mutation paths (Open/Create) use the normal locking.
// already serialize via full Lock. The O(maxEvictBatch) walk is negligible vs. deletes.
func (m *MemoryFS) EvictLRU(bytesNeeded uint) uint {
	m.mu.Lock()
	var toEvict []string
	need := int64(bytesNeeded)
	cur := m.size
	for cur > m.capacity-need && m.LRU.Len() > 0 && len(toEvict) < maxEvictBatch {
		elem := m.LRU.Back()
		if elem == nil {
			break
		}
		fi := elem.Value.(*types.FileInfo)
		key := fi.Key
		m.LRU.Remove(key) // actually remove during collection so Back() advances to distinct items
		toEvict = append(toEvict, key)
		cur -= fi.Size // local estimate; real size updated in W phase
	}
	m.mu.Unlock()

	if len(toEvict) == 0 {
		return 0
	}

	m.mu.Lock()
	var evicted uint
	for _, key := range toEvict {
		if fi, exists := m.info[key]; exists {
			m.LRU.Remove(key)
			delete(m.info, key)
			delete(m.data, key)
			m.size -= fi.Size
			evicted += uint(fi.Size)
			shardIndex := locks.GetShardIndex(key)
			m.keyLocks[shardIndex].Delete(key)
		}
	}
	m.mu.Unlock()
	return evicted
}

// EvictBySize evicts files by size (ascending = smallest first, descending = largest first)
// Collect scalar snapshot (key+size) under RLock (no shared *FileInfo pointers),
// sort on copy, brief WLock with live re-fetch for size subtract (fixes data race + accounting drift).
type evictCandidate struct {
	key  string
	size int64
}

func (m *MemoryFS) EvictBySize(bytesNeeded uint, ascending bool) uint {
	m.mu.RLock()
	var candidates []evictCandidate
	for key, fi := range m.info {
		candidates = append(candidates, evictCandidate{key: key, size: fi.Size})
		if len(candidates) >= maxEvictBatch {
			break
		}
	}
	m.mu.RUnlock()

	if len(candidates) == 0 {
		return 0
	}

	sort.Slice(candidates, func(i, j int) bool {
		if ascending {
			return candidates[i].size < candidates[j].size
		}
		return candidates[i].size > candidates[j].size
	})

	m.mu.Lock()
	var evicted uint
	for _, c := range candidates {
		if m.size <= m.capacity-int64(bytesNeeded) {
			break
		}
		key := c.key
		if liveFi, exists := m.info[key]; exists {
			m.LRU.Remove(key)
			delete(m.info, key)
			delete(m.data, key)
			m.size -= liveFi.Size
			evicted += uint(liveFi.Size)
			shardIndex := locks.GetShardIndex(key)
			m.keyLocks[shardIndex].Delete(key)
		}
	}
	m.mu.Unlock()
	return evicted
}

// EvictFIFO evicts files using FIFO (oldest creation time first)
// Collect scalar snapshot (key+ctime) under RLock, sort on copy, W phase with live re-fetch.
func (m *MemoryFS) EvictFIFO(bytesNeeded uint) uint {
	m.mu.RLock()
	var candidates []struct {
		key   string
		cTime time.Time
	}
	for key, fi := range m.info {
		candidates = append(candidates, struct {
			key   string
			cTime time.Time
		}{key: key, cTime: fi.CTime})
		if len(candidates) >= maxEvictBatch {
			break
		}
	}
	m.mu.RUnlock()

	if len(candidates) == 0 {
		return 0
	}
	sort.Slice(candidates, func(i, j int) bool {
		return candidates[i].cTime.Before(candidates[j].cTime)
	})

	m.mu.Lock()
	var evicted uint
	for _, c := range candidates {
		if m.size <= m.capacity-int64(bytesNeeded) {
			break
		}
		key := c.key
		if liveFi, exists := m.info[key]; exists {
			m.LRU.Remove(key)
			delete(m.info, key)
			delete(m.data, key)
			m.size -= liveFi.Size
			evicted += uint(liveFi.Size)
			shardIndex := locks.GetShardIndex(key)
			m.keyLocks[shardIndex].Delete(key)
		}
	}
	m.mu.Unlock()
	return evicted
}

// EvictLFU evicts least frequently used files first (by AccessCount ascending).
// Ties broken by ATime (older first). Uses scalar snapshot under RLock + live re-fetch under WLock.
func (m *MemoryFS) EvictLFU(bytesNeeded uint) uint {
	m.mu.RLock()
	var candidates []struct {
		key         string
		accessCount int
		aTime       time.Time
	}
	for key, fi := range m.info {
		candidates = append(candidates, struct {
			key         string
			accessCount int
			aTime       time.Time
		}{key: key, accessCount: fi.AccessCount, aTime: fi.ATime})
		if len(candidates) >= maxEvictBatch {
			break
		}
	}
	m.mu.RUnlock()

	if len(candidates) == 0 {
		return 0
	}
	sort.Slice(candidates, func(i, j int) bool {
		if candidates[i].accessCount != candidates[j].accessCount {
			return candidates[i].accessCount < candidates[j].accessCount
		}
		return candidates[i].aTime.Before(candidates[j].aTime)
	})

	m.mu.Lock()
	var evicted uint
	for _, c := range candidates {
		if m.size <= m.capacity-int64(bytesNeeded) {
			break
		}
		key := c.key
		if liveFi, exists := m.info[key]; exists {
			m.LRU.Remove(key)
			delete(m.info, key)
			delete(m.data, key)
			m.size -= liveFi.Size
			evicted += uint(liveFi.Size)
			shardIndex := locks.GetShardIndex(key)
			m.keyLocks[shardIndex].Delete(key)
		}
	}
	m.mu.Unlock()
	return evicted
}

// EvictHybrid evicts using time-decayed score (recency + frequency from GetTimeDecayedScore; lower value first).
// This makes "hybrid" a meaningful size + recency + frequency policy.
// Snapshot fields under RLock,
// compute score from snapshot in sort (avoids live pointer + time race post-unlock).
func (m *MemoryFS) EvictHybrid(bytesNeeded uint) uint {
	m.mu.RLock()
	var candidates []struct {
		key         string
		accessCount int
		aTime       time.Time
	}
	for key, fi := range m.info {
		candidates = append(candidates, struct {
			key         string
			accessCount int
			aTime       time.Time
		}{key: key, accessCount: fi.AccessCount, aTime: fi.ATime})
		if len(candidates) >= maxEvictBatch {
			break
		}
	}
	m.mu.RUnlock()

	if len(candidates) == 0 {
		return 0
	}
	sort.Slice(candidates, func(i, j int) bool {
		// Compute from snapshot scalars using shared DecayedScore (single source of truth).
		scoreI := types.DecayedScore(candidates[i].aTime, candidates[i].accessCount)
		scoreJ := types.DecayedScore(candidates[j].aTime, candidates[j].accessCount)
		return scoreI < scoreJ
	})

	m.mu.Lock()
	var evicted uint
	for _, c := range candidates {
		if m.size <= m.capacity-int64(bytesNeeded) {
			break
		}
		key := c.key
		if liveFi, exists := m.info[key]; exists {
			m.LRU.Remove(key)
			delete(m.info, key)
			delete(m.data, key)
			m.size -= liveFi.Size
			evicted += uint(liveFi.Size)
			shardIndex := locks.GetShardIndex(key)
			m.keyLocks[shardIndex].Delete(key)
		}
	}
	m.mu.Unlock()
	return evicted
}