package api

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"time"
)

// wrapWithMetrics wraps a handler with Prometheus metrics tracking
func (s *Server) wrapWithMetrics(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.promMetrics == nil {
			next.ServeHTTP(w, r)
			return
		}

		start := time.Now()

		// Track HTTP request
		method := r.Method
		endpoint := r.URL.Path

		// Wrap response writer to capture status code
		ww := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}

		// Serve request
		next.ServeHTTP(ww, r)

		// Record metrics
		duration := time.Since(start)
		statusStr := http.StatusText(ww.statusCode)

		s.promMetrics.IncHTTPRequests(method, endpoint, statusStr)
		s.promMetrics.ObserveHTTPDuration(method, endpoint, duration)
	})
}

// responseWriter wraps http.ResponseWriter to capture status code
type responseWriter struct {
	http.ResponseWriter
	statusCode int
}

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

func (rw *responseWriter) Flush() {
	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	h, ok := rw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("websocket: response does not implement http.Hijacker")
	}
	return h.Hijack()
}