// Package bufferedfilewriter provides a writer that buffers data in memory until a threshold is exceeded at
// which point it switches to writing to a temporary file.
package bufferedfilewriter

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer"
	"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/pool"
	"github.com/trufflesecurity/trufflehog/v3/pkg/cleantemp"
)

type bufferedFileWriterMetrics struct{}

func (bufferedFileWriterMetrics) recordDataProcessed(size uint64, dur time.Duration) {
	totalWriteSize.Add(float64(size))
	totalWriteDuration.Add(float64(dur.Microseconds()))
}

func (bufferedFileWriterMetrics) recordDiskWrite(size int64) {
	diskWriteCount.Inc()
	fileSizeHistogram.Observe(float64(size))
}

type PoolSize int

const (
	Default PoolSize = iota
	Large
)

const (
	defaultBufferSize = 1 << 12 // 4KB
	largeBufferSize   = 1 << 16 // 64KB
)

func init() {
	defaultBufferPool = pool.NewBufferPool(defaultBufferSize)
	largeBufferPool = pool.NewBufferPool(largeBufferSize)
}

// Different buffer pools for different buffer sizes.
// This allows for more efficient memory management based on the size of the data being written.
var (
	defaultBufferPool *pool.Pool
	largeBufferPool   *pool.Pool
)

// state represents the current mode of BufferedFileWriter.
type state uint8

const (
	// writeOnly indicates the BufferedFileWriter is in write-only mode.
	writeOnly state = iota
	// readOnly indicates the BufferedFileWriter has been closed and is in read-only mode.
	readOnly
)

// BufferedFileWriter manages a buffer for writing data, flushing to a file when a threshold is exceeded.
// It supports either write-only or read-only mode, indicated by its state.
type BufferedFileWriter struct {
	threshold uint64 // Threshold for switching to file writing.
	size      uint64 // Total size of the data written.

	bufPool  *pool.Pool     // Pool for storing buffers for reuse.
	buf      *buffer.Buffer // Buffer for storing data under the threshold in memory.
	filename string         // Name of the temporary file.
	file     *os.File       // File for storing data over the threshold.

	state state // Current state of the writer. (writeOnly or readOnly)

	metrics bufferedFileWriterMetrics
}

// Option is a function that modifies a BufferedFileWriter.
type Option func(*BufferedFileWriter)

// WithThreshold sets the threshold for switching to file writing.
func WithThreshold(threshold uint64) Option {
	return func(w *BufferedFileWriter) { w.threshold = threshold }
}

// WithBufferSize sets the buffer size for the BufferedFileWriter.
func WithBufferSize(size PoolSize) Option {
	return func(w *BufferedFileWriter) {
		switch size {
		case Default:
			w.bufPool = defaultBufferPool
		case Large:
			w.bufPool = largeBufferPool
		default:
			w.bufPool = defaultBufferPool
		}
	}
}

const defaultThreshold = 10 * 1024 * 1024 // 10MB
// New creates a new BufferedFileWriter with the given options.
func New(opts ...Option) *BufferedFileWriter {
	w := &BufferedFileWriter{
		threshold: defaultThreshold,
		state:     writeOnly,
	}

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

	if w.bufPool == nil {
		w.bufPool = defaultBufferPool
	}

	return w
}

// Len returns the number of bytes written to the buffer or file.
func (w *BufferedFileWriter) Len() int { return int(w.size) }

// String returns all the data written to the buffer or file as a string or an empty string if there is an error.
func (w *BufferedFileWriter) String() (string, error) {
	if w.file == nil {
		return w.buf.String(), nil
	}

	// Data is in a file, read from the file.
	file, err := os.Open(w.filename)
	if err != nil {
		return "", fmt.Errorf("failed to open file: %w", err)
	}
	defer file.Close()

	var buf bytes.Buffer
	// Read the file contents into the buffer.
	if _, err := io.CopyBuffer(&buf, file, nil); err != nil {
		return "", fmt.Errorf("failed to read file contents: %w", err)
	}

	// Append buffer data, if any, to the end of the file contents.
	if _, err := w.buf.WriteTo(&buf); err != nil {
		return "", err
	}

	return buf.String(), nil
}

// Write writes data to the buffer or a file, depending on the size.
func (w *BufferedFileWriter) Write(data []byte) (int, error) {
	if w.state != writeOnly {
		return 0, fmt.Errorf("BufferedFileWriter must be in write-only mode to write")
	}

	if w.buf == nil {
		w.buf = w.bufPool.Get()
		if w.buf == nil {
			w.buf = buffer.NewBuffer()
		}
	}

	size := uint64(len(data))
	bufferLength := w.buf.Len()
	start := time.Now()
	defer func(start time.Time) {
		w.metrics.recordDataProcessed(size, time.Since(start))
		w.size += size
	}(start)

	totalSizeNeeded := uint64(bufferLength) + size
	if totalSizeNeeded <= w.threshold {
		return w.buf.Write(data)
	}

	// Switch to file writing if threshold is exceeded.
	// This helps in managing memory efficiently for large content.
	if w.file == nil {
		file, err := os.CreateTemp(os.TempDir(), cleantemp.MkFilename())
		if err != nil {
			return 0, err
		}

		w.filename = file.Name()
		w.file = file

		// Transfer existing data in buffer to the file, then clear the buffer.
		// This ensures all the data is in one place - either entirely in the buffer or the file.
		if bufferLength > 0 {
			if _, err := w.buf.WriteTo(w.file); err != nil {
				if err := os.RemoveAll(w.filename); err != nil {
					return 0, fmt.Errorf("failed to remove file: %w", err)
				}
				return 0, err
			}
		}
	}

	// Write any remaining data in the buffer to the file before writing new data.
	if w.buf.Len() > 0 {
		if _, err := w.buf.WriteTo(w.file); err != nil {
			return 0, fmt.Errorf("error flushing buffer to file: %w", err)
		}
	}

	n, err := w.file.Write(data)
	if err != nil {
		return n, err
	}
	w.metrics.recordDiskWrite(int64(n))

	return n, nil
}

// ReadFrom reads data from the provided reader and writes it to the buffer or file, depending on the size.
// This method satisfies the io.ReaderFrom interface, allowing it to be used with standard library functions
// like io.Copy and io.CopyBuffer.
//
// By implementing this method, BufferedFileWriter can leverage optimized data transfer mechanisms provided
// by the standard library. For example, when using io.Copy with a BufferedFileWriter, the copy operation
// will be delegated to the ReadFrom method, avoiding the potentially non-optimized default approach.
//
// This is particularly useful when creating a new BufferedFileWriter from an io.Reader using the NewFromReader
// function. By leveraging the ReadFrom method, data can be efficiently transferred from the reader to
// the BufferedFileWriter.
func (w *BufferedFileWriter) ReadFrom(reader io.Reader) (int64, error) {
	if w.state != writeOnly {
		return 0, fmt.Errorf("BufferedFileWriter must be in write-only mode to write")
	}

	var totalBytesRead int64
	const bufferSize = 1 << 16 // 64KB
	buf := make([]byte, bufferSize)

	for {
		n, err := reader.Read(buf)
		if err != nil && !errors.Is(err, io.EOF) {
			return totalBytesRead, err
		}

		if n > 0 {
			written, err := w.Write(buf[:n])
			if err != nil {
				return totalBytesRead, err
			}
			totalBytesRead += int64(written)
		}

		if errors.Is(err, io.EOF) {
			break
		}
	}

	return totalBytesRead, nil
}

// CloseForWriting flushes any remaining data in the buffer to the file, closes the file if created,
// and transitions the BufferedFileWriter to read-only mode.
func (w *BufferedFileWriter) CloseForWriting() error {
	defer func() { w.state = readOnly }()
	if w.file == nil {
		return nil
	}

	// Return the buffer to the pool since the contents have been written to the file and
	// the writer is transitioning to read-only mode.
	defer w.bufPool.Put(w.buf)

	if w.buf.Len() > 0 {
		if _, err := w.buf.WriteTo(w.file); err != nil {
			return err
		}
	}
	return w.file.Close()
}

// ReadCloser returns an io.ReadCloser to read the written content.
// If the content is stored in a file, it opens the file and returns a file reader.
// If the content is stored in memory, it returns a custom reader that handles returning the buffer to the pool.
// The caller should call Close() on the returned io.Reader when done to ensure resources are properly released.
// This method can only be used when the BufferedFileWriter is in read-only mode.
func (w *BufferedFileWriter) ReadCloser() (io.ReadCloser, error) {
	if w.state != readOnly {
		return nil, fmt.Errorf("BufferedFileWriter must be in read-only mode to read")
	}

	if w.file != nil {
		// Data is in a file, read from the file.
		file, err := os.Open(w.filename)
		if err != nil {
			return nil, err
		}
		return newAutoDeletingFileReader(file), nil
	}

	if w.buf == nil {
		return nil, nil
	}

	// Data is in memory.
	return buffer.ReadCloser(w.buf.Bytes(), func() { w.bufPool.Put(w.buf) }), nil
}

// autoDeletingFileReader wraps an *os.File and deletes the file on Close.
type autoDeletingFileReader struct{ *os.File }

// newAutoDeletingFileReader creates a new autoDeletingFileReader.
func newAutoDeletingFileReader(file *os.File) *autoDeletingFileReader {
	return &autoDeletingFileReader{File: file}
}

// Close implements the io.Closer interface, deletes the file after closing.
func (r *autoDeletingFileReader) Close() error {
	defer os.Remove(r.Name()) // Delete the file after closing
	return r.File.Close()
}
