package middlewares

import (
	"net/http"

	"code.aides.space/aides-infra/aides-repo-api/internal/logger"
)

// CustomResponseWriter wraps http.ResponseWriter to capture status code and response size
type CustomResponseWriter struct {
	http.ResponseWriter
	StatusCode   int
	BytesWritten int
}

// WriteHeader captures the status code and writes it to the response
func (w *CustomResponseWriter) WriteHeader(statusCode int) {
	w.StatusCode = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

// Write captures the size of the response body
func (w *CustomResponseWriter) Write(data []byte) (int, error) {
	bytesWritten, err := w.ResponseWriter.Write(data)
	w.BytesWritten += bytesWritten
	return bytesWritten, err
}

func LoggerMiddleware(logger logger.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			customWriter := &CustomResponseWriter{
				ResponseWriter: w,
				StatusCode:     http.StatusOK, // Default status if WriteHeader is not called
			}
			next.ServeHTTP(customWriter, r)
			logger.Info("request", map[string]interface{}{
				"path":          r.URL.Path,
				"method":        r.Method,
				"ip":            r.RemoteAddr,
				"status_code":   customWriter.StatusCode,
				"response_size": customWriter.BytesWritten,
			})
		})
	}
}