package proxy

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

// statusRecorder wraps http.ResponseWriter to capture status code and
// bytes written for logging/metrics, while still supporting hijacking
// (required for WebSocket upgrades) and flushing (required for streaming
// responses / SSE).
type statusRecorder struct {
	http.ResponseWriter
	status      int
	bytes       int64
	wroteHeader bool
}

func (w *statusRecorder) WriteHeader(code int) {
	if w.wroteHeader {
		return
	}
	w.status = code
	w.wroteHeader = true
	w.ResponseWriter.WriteHeader(code)
}

func (w *statusRecorder) Write(b []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	n, err := w.ResponseWriter.Write(b)
	w.bytes += int64(n)
	return n, err
}

func (w *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hj, ok := w.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("underlying ResponseWriter does not support hijacking")
	}
	// A hijacked WebSocket connection never goes through WriteHeader, so
	// treat it as a successful 101 for logging purposes.
	w.status = http.StatusSwitchingProtocols
	w.wroteHeader = true
	return hj.Hijack()
}

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