package middleware

import (
	"net/http"
	"time"

	"github.com/rs/zerolog"
	"github.com/rs/zerolog/hlog"
)

// responseWriter wraps http.ResponseWriter to capture status code and bytes written.
type responseWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

func (rw *responseWriter) WriteHeader(code int) {
	rw.status = code
	rw.ResponseWriter.WriteHeader(code)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	n, err := rw.ResponseWriter.Write(b)
	rw.bytes += n
	return n, err
}

// Logger returns a zerolog structured logging middleware. It logs each request
// with method, path, status, latency, remote IP, and request ID (if set by
// hlog.RequestIDHandler upstream).
func Logger(log zerolog.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()

			rw := &responseWriter{ResponseWriter: w, status: http.StatusOK}
			next.ServeHTTP(rw, r)

			latency := time.Since(start)

			event := log.Info()

			// Attach request ID if hlog placed one in context.
			if id, ok := hlog.IDFromRequest(r); ok {
				event = event.Str("req_id", id.String())
			}

			event.
				Str("method", r.Method).
				Str("path", r.URL.Path).
				Str("query", r.URL.RawQuery).
				Int("status", rw.status).
				Int("bytes", rw.bytes).
				Dur("latency_ms", latency).
				Str("remote_ip", realIP(r)).
				Str("user_agent", r.UserAgent()).
				Msg("request")
		})
	}
}

// realIP returns the originating IP, preferring X-Forwarded-For / X-Real-IP
// headers set by a trusted reverse proxy.
func realIP(r *http.Request) string {
	if ip := r.Header.Get("X-Real-IP"); ip != "" {
		return ip
	}
	if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
		// X-Forwarded-For may be a comma-separated list; take the first entry.
		for i := 0; i < len(ip); i++ {
			if ip[i] == ',' {
				return ip[:i]
			}
		}
		return ip
	}
	return r.RemoteAddr
}