package audit

import (
	"context"
	"net/http"
)

type bodyWrapperContextKey struct{}

func bodyWrapperFromRequest(r *http.Request) *BodyWrapper {
	if r == nil {
		return nil
	}

	wrapper, ok := r.Context().Value(bodyWrapperContextKey{}).(*BodyWrapper)
	if !ok {
		return nil
	}
	return wrapper
}

// AttachBodyCapture wraps r.Body to capture request body bytes for audit logging.
// It is safe to call even when audit is disabled or when request body logging is disabled.
func (m *Manager) AttachBodyCapture(r *http.Request) *http.Request {
	if m == nil || r == nil {
		return r
	}
	// Avoid double-wrapping.
	if bodyWrapperFromRequest(r) != nil {
		return r
	}

	enabled, maxBodySize := m.GetAuditSnapshot()
	if !enabled || maxBodySize <= 0 || r.Body == nil {
		return r
	}

	wrapper := NewBodyWrapper(r.Body, maxBodySize)
	ctx := context.WithValue(r.Context(), bodyWrapperContextKey{}, wrapper)

	updated := r.WithContext(ctx)
	updated.Body = wrapper
	return updated
}

// RecordRequestAutoBody records an HTTP request to the audit log.
// If AttachBodyCapture() was used, it automatically includes the captured request body
// when policy allows it (e.g., status code < 400).
func (m *Manager) RecordRequestAutoBody(r *http.Request, opts RequestRecordOptions) {
	if m == nil || r == nil {
		return
	}

	// Defensive: never log request bodies for error responses.
	if opts.StatusCode >= http.StatusBadRequest {
		opts.Body = nil
		opts.BodyTruncated = false
		m.RecordRequest(r, opts)
		return
	}

	if wrapper := bodyWrapperFromRequest(r); wrapper != nil {
		opts.Body = wrapper.AuditBytes()
		opts.BodyTruncated = wrapper.Truncated()
	}

	m.RecordRequest(r, opts)
}