// Package wshttp contains a basic implementation of the interfaces
// 'http.ResponseWriter' and 'http.Hijacker' centered on websocket usage,
// specially the handshake response.
package wshttp

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
)

// Response implements 'http.ResponseWriter' and 'http.Hijacker'.
type Response struct {
	Res       http.Response
	Con       net.Conn
	hasStatus bool
	hijacked  bool
}

// Header implements 'http.ResponseWriter' strictly for websocket connections.
func (r *Response) Header() http.Header {
	if r == nil {
		return nil
	}
	if r.Res.Header == nil {
		r.Res.Header = make(map[string][]string)
	}
	return r.Res.Header
}

// WriteHeader implements 'http.ResponseWriter' strictly for websocket
// connections.
func (r *Response) WriteHeader(code int) {
	if r == nil {
		return
	}

	if r.hasStatus {
		return
	}
	r.hasStatus = true

	r.Res.StatusCode = code
	r.Res.Status = http.StatusText(code)
}

// Write implements 'http.ResponseWriter' strictly for websocket connections.
func (r *Response) Write(data []byte) (int, error) {
	if r == nil {
		return 0, fmt.Errorf("nil response")
	}
	if r.hijacked {
		return 0, http.ErrHijacked
	}

	if !r.hasStatus {
		r.Res.StatusCode = http.StatusOK
		r.Res.Status = http.StatusText(http.StatusOK)
	}

	if r.Res.Body == nil {
		r.Res.Body = io.NopCloser(bytes.NewBuffer(data))
	}
	var dump, err = httputil.DumpResponse(&r.Res, true)
	if err != nil {
		return 0, err
	}

	return r.Con.Write(dump)
}

// Hijack implements 'http.Hijacker'. The response body is drained, as it's not
// used in websocket handshakes.
func (r *Response) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if r == nil {
		return nil, nil, fmt.Errorf("nil response")
	}
	if r.hijacked {
		return nil, nil, http.ErrHijacked
	}
	var err error

	r.hijacked = true
	if r.Res.Body != nil {
		if _, err = io.Copy(io.Discard, r.Res.Body); err != nil {
			return nil, nil, fmt.Errorf("unable to drain response body: %v", err)
		}
	}

	return r.Con, bufio.NewReadWriter(bufio.NewReader(r.Con), bufio.NewWriter(r.Con)), nil
}

// NewResponse instantiates an HTTP response for websocket handshake requests.
func NewResponse(con net.Conn) *Response {
	return &Response{
		Con: con,
		Res: http.Response{Header: make(map[string][]string)},
	}
}