Pad/internal/editor/chunked_buffer.go

package editor

import (
	"bytes"
	"fmt"
	"sort"

	"pad/internal/io/pool"
	"pad/internal/io/pool/types"
	"pad/internal/ui"
)

const (
	DefaultChunkSize = 64 * 1024 // 64 KB
)

// ChunkedBuffer provides chunked access to a file's content.
// Only chunks near the cursor or viewport are kept in memory.
type ChunkedBuffer struct {
	filename   string
	chunkSize  int // e.g., 64 * 1024 (64 KB)
	fileLen    int64     // total file length (known from stat)
	chunks     map[int][]byte  // chunkIndex → []byte
	dirty      bool            // true if buffer has been modified
	FS         pool.FileSystem // filesystem for reading chunks
	basePath   string          // base path for file resolution

	// line index is built asynchronously
	LineIndex *types.LineIndex

	// workerPool is set by the logic goroutine after the buffer is created.
	// Used for async prefetch of chunks.
	workerPool *pool.WorkerPool

	// lastPrefetchedChunk tracks the last chunk that was prefetched,
	// so we can avoid redundant prefetches on the same position.
	lastPrefetchedChunk int

	// dirtyChunks tracks which individual chunks have been modified
	// since they were last persisted to disk. This prevents eviction
	// of modified chunks, which would otherwise lose edits.
	dirtyChunks map[int]bool
}

// NewChunkedBuffer creates a new ChunkedBuffer.
func NewChunkedBuffer(filename string, chunkSize int, fs pool.FileSystem, basePath string) *ChunkedBuffer {
	if chunkSize <= 0 {
		chunkSize = DefaultChunkSize
	}
	return &ChunkedBuffer{
		filename:    filename,
		chunkSize:   chunkSize,
		chunks:      make(map[int][]byte),
		dirtyChunks: make(map[int]bool),
		FS:          fs,
		basePath:    basePath,
	}
}

// SetFileSize sets the total file length.
func (cb *ChunkedBuffer) SetFileSize(length int64) {
	cb.fileLen = length
}

// FileLen returns the total file length.
func (cb *ChunkedBuffer) FileLen() int64 {
	return cb.fileLen
}

// Filename returns the filename.
func (cb *ChunkedBuffer) Filename() string {
	return cb.filename
}

// ChunkSize returns the chunk size.
func (cb *ChunkedBuffer) ChunkSize() int {
	return cb.chunkSize
}

// Content returns the bytes in [start, end) from the chunked buffer.
// It loads missing chunks on demand.
func (cb *ChunkedBuffer) Content(start, end int) string {
	if cb.fileLen == 0 {
		return ""
	}
	// Clamp range to file bounds
	if start < 0 {
		start = 0
	}
	if end > int(cb.fileLen) {
		end = int(cb.fileLen)
	}
	if start >= end {
		return ""
	}

	startChunk := start / cb.chunkSize
	endChunk := (end - 1) / cb.chunkSize

	var buf bytes.Buffer
	for i := startChunk; i <= endChunk; i++ {
		chunk, ok := cb.chunks[i]
		if !ok {
			// If chunk is not loaded, load it. This is a blocking call.
			// In a real app, this might be async or a fallback.
			loadedChunk, err := cb.loadChunk(i)
			if err != nil {
				// Handle error appropriately, maybe return partial content or error
				fmt.Printf("Error loading chunk %d for file %s: %v\n", i, cb.filename, err)
				continue // Skip this chunk on error
			}
			chunk = loadedChunk
		}

		chunkStart := i * cb.chunkSize
		chunkEnd := chunkStart + len(chunk)

		segStart := max(start, chunkStart) - chunkStart
		segEnd := min(end, chunkEnd) - chunkStart

		if segStart < segEnd {
			buf.Write(chunk[segStart:segEnd])
		}
	}
	return buf.String()
}

// FullContent reconstructs the entire file content from loaded or re-read chunks.
// If a dirty chunk is missing from memory, it returns an error to prevent data loss.
func (cb *ChunkedBuffer) FullContent() (string, error) {
	if cb.fileLen == 0 {
		return "", nil
	}
	numChunks := int((cb.fileLen + int64(cb.chunkSize) - 1) / int64(cb.chunkSize))
	var buf bytes.Buffer
	for i := 0; i < numChunks; i++ {
		var chunk []byte
		var err error
		if loadedChunk, ok := cb.chunks[i]; ok {
			chunk = loadedChunk
		} else {
			// If chunk is dirty but not in memory, we have a problem.
			if cb.dirtyChunks[i] {
				return "", fmt.Errorf("critical error: dirty chunk %d is missing from memory", i)
			}
			// Re-read from disk if not in memory
			chunk, err = cb.loadChunk(i)
			if err != nil {
				fmt.Printf("Error re-reading chunk %d for file %s during FullContent: %v\n", i, cb.filename, err)
				continue // Skip this chunk on error
			}
		}
		buf.Write(chunk)
	}
	return buf.String(), nil
}

// loadChunk reads a specific chunk from disk and returns its content.
// It also stores the chunk in the 'chunks' map.
// Uses FS.ReadFileAt to read only the chunk range (not the entire file).
func (cb *ChunkedBuffer) loadChunk(idx int) ([]byte, error) {
	start := idx * cb.chunkSize
	chunk, err := cb.FS.ReadFileAt(cb.filename, start, cb.chunkSize)
	if err != nil {
		return nil, fmt.Errorf("failed to read chunk %d: %w", idx, err)
	}
	cb.chunks[idx] = chunk // Store the loaded chunk
	return chunk, nil
}

// LoadChunk explicitly loads a chunk and prefetch adjacent ones.
func (cb *ChunkedBuffer) LoadChunk(idx int) {
	if _, ok := cb.chunks[idx]; !ok {
		_, err := cb.loadChunk(idx)
		if err != nil {
			fmt.Printf("Error loading chunk %d: %v\n", idx, err)
		}
	}
	// Prefetch adjacent chunks
	cb.Prefetch(idx, 1)
}

// SetWorkerPool sets the worker pool for async chunk loading.
func (cb *ChunkedBuffer) SetWorkerPool(wp *pool.WorkerPool) {
	cb.workerPool = wp
}

// LastPrefetchedChunk returns the last chunk that was prefetched.
// Used by EditorLayout to avoid redundant prefetches.
func (cb *ChunkedBuffer) LastPrefetchedChunk() int {
	return cb.lastPrefetchedChunk
}

// Prefetch loads adjacent chunks for smooth scrolling.
// radius is the number of chunks to load on each side.
// Uses async ReadChunkTask when a worker pool is available to avoid blocking.
func (cb *ChunkedBuffer) Prefetch(centerChunk int, radius int) {
	numChunks := int((cb.fileLen + int64(cb.chunkSize) - 1) / int64(cb.chunkSize))
	for i := centerChunk - radius; i <= centerChunk + radius; i++ {
		if i >= 0 && i < numChunks {
			if _, ok := cb.chunks[i]; !ok {
				if cb.workerPool != nil {
					// Dispatch async read chunk task
					cb.workerPool.DispatchNonBlocking(
						pool.NewReadChunkTask(cb.filename, i, cb.FS),
					)
				} else {
					// Fallback: synchronous load (should only happen during testing)
					_, err := cb.loadChunk(i)
					if err != nil {
						fmt.Printf("Error prefetching chunk %d: %v\n", i, err)
					}
				}
			}
		}
	}
	cb.lastPrefetchedChunk = centerChunk
}

// EvictFarChunks removes chunks that are too far from the cursor.
// radius is the number of chunks to keep around the cursor.
// Dirty chunks (modified since the last disk write) are NEVER evicted,
// even if they are far from the cursor, to prevent data loss.
func (cb *ChunkedBuffer) EvictFarChunks(cursorPos int, radius int) {
	if cb.fileLen == 0 {
		return
	}
	cursorChunk := cursorPos / cb.chunkSize
	numChunks := int((cb.fileLen + int64(cb.chunkSize) - 1) / int64(cb.chunkSize))

	for i := range cb.chunks {
		// NEVER evict dirty chunks — they contain unsaved modifications
		if cb.dirtyChunks[i] {
			continue
		}
		if i < cursorChunk-radius || i > cursorChunk+radius {
			// Check if chunk index is within valid range before deleting
			if i >= 0 && i < numChunks {
				delete(cb.chunks, i)
			}
		}
	}
}

// Insert inserts text at a given position.
func (cb *ChunkedBuffer) Insert(pos int, text string) {
	if len(text) == 0 {
		return
	}
	// Ensure the chunk containing pos is loaded
	chunkIdx := pos / cb.chunkSize
	// Ensure we don't try to insert beyond the current fileLen if it's not a new file
	if pos > int(cb.fileLen) && cb.fileLen > 0 {
		pos = int(cb.fileLen) // clamp insertion point to end of file
	}
	if pos < 0 {
		pos = 0
	}

	// Load chunk if it doesn't exist, or if pos is at the very beginning of a non-loaded chunk
	if _, ok := cb.chunks[chunkIdx]; !ok {
		// If inserting at the start of a chunk, we need to load it.
		// If inserting past the end of the file, we might create new chunks.
		// For now, assume loadChunk handles cases where pos is beyond current fileLen by reading up to fileLen.
		_, err := cb.loadChunk(chunkIdx) // This is blocking and might be problematic
		if err != nil {
			fmt.Printf("Error loading chunk %d for insert: %v\n", chunkIdx, err)
			return
		}
	}

	chunk := cb.chunks[chunkIdx]
	offsetInChunk := pos - chunkIdx*cb.chunkSize

	// Ensure offsetInChunk is valid for the loaded chunk
	if offsetInChunk > len(chunk) {
		// This can happen if we are inserting past the end of the loaded chunk,
		// which might be due to fileLen not being updated or insertion into new territory.
		// For now, we'll pad the chunk if needed. This needs more robust handling.
		padding := make([]byte, offsetInChunk-len(chunk))
		chunk = append(chunk, padding...)
	}

	// Insert into the chunk
	newChunk := make([]byte, len(chunk)+len(text))
	copy(newChunk, chunk[:offsetInChunk])
	copy(newChunk[offsetInChunk:], text)
	copy(newChunk[offsetInChunk+len(text):], chunk[offsetInChunk:])
	cb.chunks[chunkIdx] = newChunk

	// Update file length if insertion extends beyond current length
	if pos+len(text) > int(cb.fileLen) {
		cb.fileLen = int64(pos + len(text))
	}

	cb.dirty = true
	cb.dirtyChunks[chunkIdx] = true

	// If insertion spans chunk boundary, it might require merging chunks.
	// This is complex and might involve resizing subsequent chunks and potentially
	// re-reading them. For now, we defer complex merge logic.
	// The plan mentions maybeMergeChunk, which would handle this.
	// cb.maybeMergeChunk(chunkIdx)
}

// Delete deletes n bytes starting at pos.
func (cb *ChunkedBuffer) Delete(pos, n int) {
	if n <= 0 {
		return
	}
	// Clamp pos and n to valid range
	if pos < 0 {
		pos = 0
	}
	if pos >= int(cb.fileLen) {
		return // Nothing to delete
	}
	if pos+n > int(cb.fileLen) {
		n = int(cb.fileLen) - pos
	}

	startChunk := pos / cb.chunkSize
	endChunk := (pos + n - 1) / cb.chunkSize

	// Load all affected chunks
	for i := startChunk; i <= endChunk; i++ {
		if i < 0 {
			continue
		}
		if _, ok := cb.chunks[i]; !ok {
			_, err := cb.loadChunk(i)
			if err != nil {
				fmt.Printf("Error loading chunk %d for delete: %v\n", i, err)
				return // Abort delete on error
			}
		}
	}

	// Perform deletion across all affected chunks.
	// We track how many bytes remain to delete and advance pos as we go.
	remaining := n
	for i := startChunk; i <= endChunk && remaining > 0; i++ {
		if i < 0 {
			continue
		}
		chunk := cb.chunks[i]
		if chunk == nil {
			continue
		}

		currentChunkStart := i * cb.chunkSize
		effectiveOffsetInChunk := max(pos, currentChunkStart) - currentChunkStart
		effectiveDeleteEnd := min(pos+remaining, currentChunkStart+len(chunk)) - currentChunkStart

		if effectiveOffsetInChunk >= len(chunk) {
			continue // Deletion range is beyond this chunk
		}

		// Perform deletion within the chunk
		cb.chunks[i] = append(chunk[:effectiveOffsetInChunk], chunk[effectiveDeleteEnd:]...)

		// Track how many bytes we deleted from this chunk
		bytesDeleted := effectiveDeleteEnd - effectiveOffsetInChunk
		remaining -= bytesDeleted

		// Mark this chunk as dirty so it won't be evicted
		cb.dirtyChunks[i] = true
	}

	// Update file length
	cb.fileLen -= int64(n)
	if cb.fileLen < 0 {
		cb.fileLen = 0
	}

	cb.dirty = true
}


// VisibleByteRange returns the byte range [start, end) of content
// visible in the viewport, given the current scroll offset and viewport height.
func (cb *ChunkedBuffer) VisibleByteRange(scrollOffset ui.Dp, viewportHeight ui.Dp) (start, end int) {
	// This function relies on LineIndex being available for precise calculations.
	// Fallback to estimate if LineIndex is nil.
	if cb.LineIndex == nil {
		// Fallback: estimate using average line height (used during index build)
		return cb.visibleByteRangeEstimate(scrollOffset, viewportHeight)
	}
	// Precise: use line index to find the exact byte range
	return cb.visibleByteRangePrecise(scrollOffset, viewportHeight)
}

// visibleByteRangeEstimate approximates the visible byte range using
// heuristic estimates. Used when the line index is not yet available.
func (cb *ChunkedBuffer) visibleByteRangeEstimate(scrollOffset ui.Dp, viewportHeight ui.Dp) (start, end int) {
	// This is a rough estimation. A proper implementation would need actual line height.
	// For simplicity, assuming a fixed line height based on FontSize.
	// This requires access to theme or font metrics, which is not directly available here.
	// As a fallback, let's use a simplified calculation based on estimated lines and average bytes per line.

	// Use the editor's line height constant for consistency.
	lineHeight := EditorLineHeight()

	startLine := int(scrollOffset / lineHeight)
	endLine := int((scrollOffset + viewportHeight) / lineHeight)

	// Clamp line numbers to reasonable bounds
	totalLinesEstimate := 0
	if cb.fileLen > 0 {
		totalLinesEstimate = int(cb.fileLen / 50) + 1 // Rough estimate: 50 bytes per line
	}
	if startLine < 0 { startLine = 0 }
	if endLine > totalLinesEstimate { endLine = totalLinesEstimate }
	if startLine >= endLine { endLine = startLine + 1 } // Ensure at least one line is visible

	// Convert line numbers to byte offsets using the line index if available
	// If LineIndex is nil, we fall back to a very rough byte estimation.
	if cb.LineIndex != nil {
		// Use LineIndex if it exists
		if startLine < len(cb.LineIndex.Offsets) {
			start = int(cb.LineIndex.Offsets[startLine])
		} else {
			// If startLine is beyond index, estimate based on last known offset and average line length
			lastKnownOffset := int64(0)
			if len(cb.LineIndex.Offsets) > 0 {
				lastKnownOffset = int64(cb.LineIndex.Offsets[len(cb.LineIndex.Offsets)-1])
			}
			linesBeyondIndex := startLine - (len(cb.LineIndex.Offsets) -1)
			start = int(lastKnownOffset + int64(linesBeyondIndex) * 50) // Estimate
		}

		if endLine < len(cb.LineIndex.Offsets) {
			end = int(cb.LineIndex.Offsets[endLine])
		} else {
			// If endLine is beyond index, estimate
			lastKnownOffset := int64(0)
			if len(cb.LineIndex.Offsets) > 0 {
				lastKnownOffset = int64(cb.LineIndex.Offsets[len(cb.LineIndex.Offsets)-1])
			}
			linesBeyondIndex := endLine - (len(cb.LineIndex.Offsets) -1)
			end = int(lastKnownOffset + int64(linesBeyondIndex) * 50) // Estimate
		}
	} else {
		// Rough byte estimation if no LineIndex
		start = startLine * 50 // rough estimate: 50 bytes per line
		end = endLine * 50
	}

	// Clamp to file bounds
	if cb.fileLen > 0 {
		if start < 0 { start = 0 }
		if end > int(cb.fileLen) { end = int(cb.fileLen) }
		if end <= start { end = start + cb.chunkSize } // Ensure at least one chunk's worth if range is invalid
	} else {
		start = 0
		end = 0 // Empty file
	}

	return start, end
}

// visibleByteRangePrecise uses the LineIndex to find the exact byte range.
func (cb *ChunkedBuffer) visibleByteRangePrecise(scrollOffset ui.Dp, viewportHeight ui.Dp) (start, end int) {
	if cb.LineIndex == nil || len(cb.LineIndex.Offsets) == 0 {
		// Should not happen if called after LineIndex is available, but as a safeguard:
		return cb.visibleByteRangeEstimate(scrollOffset, viewportHeight)
	}

	// Use the editor's line height constant for consistency.
	lineHeight := EditorLineHeight()

	startLine := int(scrollOffset / lineHeight)
	endLine := int((scrollOffset + viewportHeight) / lineHeight)

	// Clamp line numbers to the available range in LineIndex
	if startLine < 0 {
		startLine = 0
	}
	if startLine >= len(cb.LineIndex.Offsets) {
		startLine = len(cb.LineIndex.Offsets) - 1 // Last available line
	}

	if endLine < 0 { // Should not happen with positive viewportHeight
		endLine = 0
	}
	if endLine >= len(cb.LineIndex.Offsets) {
		endLine = len(cb.LineIndex.Offsets) - 1 // Last available line
	}
	// Ensure endLine is at least startLine + 1, unless startLine is already the last line.
	if startLine < len(cb.LineIndex.Offsets)-1 && endLine <= startLine {
		endLine = startLine + 1
	}


	start = int(cb.LineIndex.Offsets[startLine])

	// The end byte offset is the start of the *next* line after the visible range.
	// If endLine is the last line in the index, the end byte offset is the file length.
	if endLine+1 < len(cb.LineIndex.Offsets) {
		end = int(cb.LineIndex.Offsets[endLine+1])
	} else {
		end = int(cb.fileLen) // Use file length as the end if it's the last line
	}

	// Ensure range is within file bounds
	if start < 0 { start = 0 }
	if end > int(cb.fileLen) { end = int(cb.fileLen) }
	if end <= start {
		// If somehow the range is invalid, return a minimal valid range,
		// e.g., start of the start line to a bit past it, or just end of file.
		if start < int(cb.fileLen) {
			end = min(start+cb.chunkSize, int(cb.fileLen)) // At least one chunk or up to file end
		} else {
			end = start // If start is already at file end, range is empty
		}
	}

	return start, end
}

// maybeMergeChunk is a placeholder for logic that consolidates chunks if they become too small
// or if edits cause fragmentation. This is complex and deferred.
func (cb *ChunkedBuffer) maybeMergeChunk(chunkIdx int) {
	// Placeholder for future implementation
	// This would involve checking chunk sizes and potentially merging adjacent chunks.
}

// UpdateLineIndexAfterEdit updates the LineIndex offsets after an insert or delete.
// offsetShift is the number of bytes inserted (positive) or deleted (negative).
// editPos is the byte position where the edit occurred.
// This is a partial update: only offsets after editPos are shifted.
func (cb *ChunkedBuffer) UpdateLineIndexAfterEdit(editPos int, offsetShift int) {
	if cb.LineIndex == nil {
		return
	}
	// Find the first offset that needs updating using binary search.
	// All offsets >= editPos need to be shifted by offsetShift.
	idx := sort.Search(len(cb.LineIndex.Offsets), func(i int) bool {
		return int(cb.LineIndex.Offsets[i]) >= editPos
	})
	for i := idx; i < len(cb.LineIndex.Offsets); i++ {
		cb.LineIndex.Offsets[i] += int32(offsetShift)
	}
	// Update the file size stamp
	cb.LineIndex.Size += int64(offsetShift)
	if cb.LineIndex.Size < 0 {
		cb.LineIndex.Size = 0
	}
}

// Helper function for max
func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// Helper function for min
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}