package hijack

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

func New(w http.ResponseWriter, commit func(http.ResponseWriter)) (http.ResponseWriter, func()) {
	responseWriter := &responseWriter{
		commit: commit,
		w:      w,
	}

	if _, ok := w.(http.Hijacker); ok {
		return &responseWriterHijacker{
			responseWriter: responseWriter,
		}, responseWriter.ensureCommitted
	} else {
		return responseWriter, responseWriter.ensureCommitted
	}
}

var _ http.ResponseWriter = new(responseWriter)
var _ http.ResponseWriter = new(responseWriterHijacker)
var _ http.Hijacker = new(responseWriterHijacker)

type responseWriter struct {
	commit    func(http.ResponseWriter)
	w         http.ResponseWriter
	hijacked  bool
	writeOnce sync.Once
}

func (rw *responseWriter) ensureCommitted() {
	rw.writeOnce.Do(func() {
		if rw.hijacked {
			return
		}
		rw.commit(rw.w)
	})
}

func (rw *responseWriter) Header() http.Header {
	return rw.w.Header()
}

func (rw *responseWriter) Write(data []byte) (int, error) {
	rw.ensureCommitted()
	return rw.w.Write(data)
}

func (rw *responseWriter) WriteHeader(status int) {
	rw.ensureCommitted()
	rw.w.WriteHeader(status)
}

type responseWriterHijacker struct {
	*responseWriter
}

func (rw *responseWriterHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	rw.hijacked = true
	return rw.w.(http.Hijacker).Hijack()
}