package handler

import (
	"go-pkg-server/logging"
	"log/slog"
	"net/http"
	"time"
)

type traceWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

var _ http.ResponseWriter = &traceWriter{}

func (t *traceWriter) Write(b []byte) (int, error) {
	if t.status == 0 {
		t.WriteHeader(http.StatusOK)
	}
	c, err := t.ResponseWriter.Write(b)
	t.bytes += c
	return c, err
}

func (tw *traceWriter) WriteHeader(statusCode int) {
	if tw.status == 0 {
		tw.status = http.StatusOK
	}
	tw.status = statusCode
	tw.ResponseWriter.WriteHeader(statusCode)
}

func New(next http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		logger := logging.FromContext(ctx)

		defer func() {
			err := recover()
			if err != nil {
				logging.RecoverLog(ctx, err)
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()

		logger = logger.With("method", r.Method,
			"host", r.Host,
			"remote_address", r.RemoteAddr,
			"path", r.RequestURI,
		)

		ctx = logging.WithLogger(ctx, logger)

		r = r.WithContext(ctx)

		tw := &traceWriter{ResponseWriter: w}
		
		start := time.Now()
		next.ServeHTTP(tw, r)
		duration := time.Since(start)

		logger.Log(ctx, slog.LevelInfo, "Request Handled",
			"protocol", r.Proto,
			"bytes_recieved", r.ContentLength,
			"status", tw.status,
			"time", duration,
			"bytes_sent", tw.bytes,
			"user_agent", r.UserAgent(),
			"referer", r.Referer(),
		)

	}
}