package audit

import (
	"bytes"
	"io"
)

// limitWriter wraps a Writer and limits the number of bytes written.
// After the limit is reached, writes are silently discarded and truncated is set.
type limitWriter struct {
	w         io.Writer
	limit     int64
	written   int64
	truncated bool
}

func newLimitWriter(w io.Writer, limit int64) *limitWriter {
	return &limitWriter{w: w, limit: limit}
}

func (l *limitWriter) Write(p []byte) (int, error) {
	if l.written >= l.limit {
		// Silently discard bytes beyond limit and mark as truncated
		l.truncated = true
		return len(p), nil
	}
	remaining := l.limit - l.written
	toWrite := p
	if int64(len(p)) > remaining {
		toWrite = p[:remaining]
		l.truncated = true
	}
	n, err := l.w.Write(toWrite)
	l.written += int64(n)
	if err != nil {
		return n, err
	}
	// Return original length since we "accepted" all bytes from TeeReader's perspective
	return len(p), nil
}

// BodyWrapper wraps an io.ReadCloser to capture the first N bytes for auditing
// while allowing the full body to be read downstream.
type BodyWrapper struct {
	original    io.ReadCloser
	buffer      *bytes.Buffer
	limitWriter *limitWriter
	tee         io.Reader
}

// NewBodyWrapper creates a new BodyWrapper that captures up to maxSize bytes
// for audit logging while passing through all bytes to downstream readers.
// Note: Truncation suffix is added by Sanitizer.SanitizeBody(), not here.
func NewBodyWrapper(body io.ReadCloser, maxSize int) *BodyWrapper {
	buf := &bytes.Buffer{}
	limitW := newLimitWriter(buf, int64(maxSize))
	return &BodyWrapper{
		original:    body,
		buffer:      buf,
		limitWriter: limitW,
		tee:         io.TeeReader(body, limitW),
	}
}

// Read reads from the underlying body while capturing bytes for audit.
func (w *BodyWrapper) Read(p []byte) (int, error) {
	return w.tee.Read(p)
}

// Close closes the underlying body.
func (w *BodyWrapper) Close() error {
	return w.original.Close()
}

// AuditBytes returns the captured bytes for audit logging.
// This should only be called after the body has been fully read or
// when you want to capture whatever has been read so far.
func (w *BodyWrapper) AuditBytes() []byte {
	return w.buffer.Bytes()
}

// Truncated returns true if the captured body was truncated due to size limit.
// Use this to determine if truncation suffix should be added to audit logs.
func (w *BodyWrapper) Truncated() bool {
	return w.limitWriter.truncated
}