package http

import (
	"bufio"
	"errors"
	"net"
	"net/http"

	"github.com/felixge/httpsnoop"
)

// Custom RW implementation
type ResponseWriter struct {
	http.ResponseWriter
	code     int
	Hijacker http.Hijacker
}

// Write status code to header
func (w *ResponseWriter) WriteHeader(statusCode int) {
	w.code = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

// GetStatusCode returns the status code that was written
func (w *ResponseWriter) GetStatusCode() int {
	return w.code
}

// Get new RW
func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
	hijacker, _ := w.(http.Hijacker)
	return &ResponseWriter{
		ResponseWriter: httpsnoop.Wrap(w, httpsnoop.Hooks{}),
		Hijacker:       hijacker,
	}
}

func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if w.Hijacker == nil {
		return nil, nil, errors.New("http.Hijacker not implemented by underlying http.ResponseWriter")
	}
	return w.Hijacker.Hijack()
}