package chitin

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"regexp"
	"runtime/pprof"
	"sync"
	"time"

	"code.justin.tv/common/chitin/internal/trace"
	"golang.org/x/net/context"
)

var (
	baseTransport *http.Transport

	httpClientProfile *pprof.Profile
)

func init() {
	baseTransport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,

			DualStack: false,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,

		TLSClientConfig:       nil,
		DisableKeepAlives:     false,
		DisableCompression:    false,
		MaxIdleConnsPerHost:   0,
		ResponseHeaderTimeout: 0,
	}

	pkgname := "code.justin.tv/common/chitin"

	// The pprof tool expects profiles to match the following regexp:
	//
	//     `\A(\w+) profile: total \d+\n\z`
	//
	// This means that the profile name can only contain the characters [0-9A-
	// Za-z_], contrary to the runtime/pprof package docs, which say:
	//
	//     The convention is to use a 'import/path.' prefix to create separate
	//     name spaces for each package.
	//
	// See https://golang.org/issue/13195
	cmdPprofName := regexp.MustCompile("[^0-9A-Za-z_]")

	httpClientProfile = pprof.NewProfile(cmdPprofName.ReplaceAllString(pkgname, "_") + "_httpclient")
}

// Client returns a Context-enhanced *http.Client.  Outbound requests will
// include the headers needed to follow user requests via Trace as they
// traverse binaries and hosts.  The Trace transaction id will be included in
// the outbound headers, and a new Trace span id will be allocated and
// included in the outbound headers as well.
//
// When the Context times out or is cancelled, any in-progress http requests
// using the returned *http.Client will be cancelled.
func Client(ctx context.Context) *http.Client {
	t, err := RoundTripper(ctx, baseTransport)
	if err != nil {
		panic(fmt.Errorf("RoundTripper must always accept a *http.Transport, and yet: %v", err))
	}

	c := &http.Client{
		Transport: t,
	}
	return c
}

// RoundTripper returns a Context-enhanced http.RoundTripper.  Outbound
// requests will include the headers needed to follow user requests via Trace
// as they traverse binaries and hosts.  The Trace transaction id will be
// included in the outbound headers, and a new Trace span id will be allocated
// and included in the outbound headers as well.
//
// When the Context times out or is cancelled, any in-progress http requests
// using the returned http.RoundTripper will be cancelled.
//
// RoundTripper will return an error if the http.RoundTripper it is provided
// does not support canceling requests (and closing idle connections).
func RoundTripper(ctx context.Context, rt http.RoundTripper) (http.RoundTripper, error) {
	if rt == nil {
		rt = baseTransport
	}

	grt, ok := rt.(goodRoundTripper)
	if !ok {
		return nil, fmt.Errorf(
			"%T does not have CancelRequest(*http.Request) and CloseIdleConnections() functions", rt)
	}

	t := &transport{
		parent: ctx,
		base:   grt,
	}
	return t, nil
}

// A goodRoundTripper allows canceling in-flight Requests.
type goodRoundTripper interface {
	http.RoundTripper
	CancelRequest(*http.Request)
	CloseIdleConnections()
}

var _ goodRoundTripper = (*http.Transport)(nil)

type transport struct {
	parent context.Context
	base   goodRoundTripper
}

func filterOutboundRequest(ctx context.Context, dst, src *http.Request, onClose func()) error {
	*dst = http.Request{
		// We scrub Method

		URL: src.URL,

		// Proto, ProtoMajor, and ProtoMinor don't apply to outbound requests.

		// We scrub Header

		// We replace Body

		ContentLength: src.ContentLength,

		TransferEncoding: src.TransferEncoding,

		// We may want to muck with Close to encourage connection churn.
		Close: src.Close,

		Host: src.Host,

		// Form, PostForm, and MultipartForm are populated by user calls

		// We don't support Trailer

		// RemoteAddr is meaningless for outbound requests.

		// RequestURI must not be set for outbound requests.

		// TLS is ignored for outbound requests.
	}

	switch src.Method {
	default:
		return fmt.Errorf("bad method: %q", src.Method)
	case "GET", "POST", "PUT", "HEAD", "DELETE", "PATCH", "OPTIONS", "":
		dst.Method = src.Method
	}

	dst.Header = make(http.Header, len(src.Header))
	for k, v := range src.Header {
		k = http.CanonicalHeaderKey(k)
		switch {
		default:
			dst.Header[k] = v
		case reservedHeader(k):
			printf(ctx, "reserved-header=%q", k)
		}
	}

	// TODO: add a test for sendfile support
	dst.Body = &clientBody{
		rc:      src.Body,
		onClose: onClose,
	}

	return nil
}

func filterInboundResponse(ctx context.Context, dst, src *http.Response, onClose func()) {
	*dst = http.Response{
		Status:     src.Status,
		StatusCode: src.StatusCode,
		Proto:      src.Proto,
		ProtoMajor: src.ProtoMajor,
		ProtoMinor: src.ProtoMinor,

		// We scrub Header

		// TODO: instrument byte count, duration, etc.
		// We replace Body

		ContentLength: src.ContentLength,

		TransferEncoding: src.TransferEncoding,

		Close: src.Close,

		// We don't support Trailer.

		// The Request must be set to the original user-provided request.

		TLS: src.TLS,
	}

	dst.Header = make(http.Header, len(src.Header))
	for k, v := range src.Header {
		k = http.CanonicalHeaderKey(k)
		switch {
		default:
			dst.Header[k] = v
		case reservedHeader(k):
			printf(ctx, "reserved-header=%q", k)
		}
	}

	dst.Body = &clientBody{
		rc:      src.Body,
		onClose: onClose,
	}
}

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	ctx := trace.NewSpan(t.parent)

	blessedReq := &http.Request{}
	err := filterOutboundRequest(ctx, blessedReq, req, func() {
		trace.SendRequestBodySent(ctx)
	})
	if err != nil {
		return nil, err
	}

	trace.AugmentHeader(ctx, blessedReq.Header)

	defer func() {
		// Mark the original body as consumed ... when the time is right.
		req.Body = blessedReq.Body
	}()

	doneHeaders := make(chan struct{})
	defer close(doneHeaders)
	doneBody := make(chan struct{})
	go t.watch(ctx, blessedReq, doneHeaders, doneBody)

	pprofKey := new(byte) // This needs to have a unique address, so it must not be size zero
	httpClientProfile.Add(pprofKey, 1)
	// It's happening!
	trace.SendRequestHeadPrepared(ctx, blessedReq)
	resp, err := t.base.RoundTrip(blessedReq)
	trace.SendResponseHeadReceived(ctx, resp, err)
	if err != nil {
		close(doneBody)
		httpClientProfile.Remove(pprofKey)
		return resp, err
	}

	blessedResp := &http.Response{}
	filterInboundResponse(ctx, blessedResp, resp, func() {
		trace.SendResponseBodyReceived(ctx)
		close(doneBody)
		httpClientProfile.Remove(pprofKey)
	})
	blessedResp.Request = req

	return blessedResp, nil
}

// watch monitors a context and cancels an outbound http request when the
// context becomes inactive.
func (t *transport) watch(ctx context.Context, req *http.Request, doneHeaders, doneBody chan struct{}) {
	// Here we track the progress of the request while monitoring the
	// context's lifecycle.
	//
	// In the first stage, we wait for the response headers to be
	// returned.  If the context becomes inactive before our client
	// receives the full response headers, we won't know how far the
	// request has progressed through the call to RoundTrip.  We'll
	// repeatedly attempt to cancel the request until the call to
	// RoundTrip returns.
	//
	// In the second stage, the response headers and body have been
	// returned to the caller.  Our caller has not yet closed the response
	// body.  Canceling at this point is much more straightforward,
	// requiring just a single cancellation attempt.
	//
	// After the response body is fully consumed, the request has
	// completed and our work is done.

	for {
		select {
		case <-ctx.Done():
			// We very much prefer not calling CancelRequest if we don't
			// have to (https://golang.org/issue/10474).
			select {
			case <-doneHeaders:
			default:
				t.base.CancelRequest(req)
				// We're racing against the RoundTripper here - it might not
				// have registered a func to run when the request is
				// cancelled. So we've got to try again and again until we
				// know that our sibling goroutine has entered and returned
				// from its call to RoundTrip.
				//
				// The value of this sleep is arbitrary - slow enough to not
				// busy-loop, fast enough to cancel the request fairly
				// promptly.
				time.Sleep(1 * time.Millisecond)
				continue
			}
		case <-doneHeaders:
		}
		break
	}
	select {
	case <-doneBody:
	case <-ctx.Done():
		// This time we only need to try once - the RoundTripper has
		// acknowledged the request and has either finished with it
		// already or will certainly cancel when we ask.

		// But, there's a logical race in the stdlib between closing the
		// TCP connection and returning it to the connection pool, thus
		// poisoning whatever unfortunate request happens to get that
		// connection next.  So, let's be a bit more certain that the body
		// hasn't already been consumed before we set loose the
		// connection-closing goroutine.
		//
		// https://golang.org/issue/10474

		select {
		case <-doneBody:
			// j/k, no need to cancel.  Apparently both cases of the
			// select were ready.
		default:
			t.base.CancelRequest(req)
		}
	}
}

func (t *transport) CancelRequest(req *http.Request) {
	t.base.CancelRequest(req)
}

func (t *transport) CloseIdleConnections() {
	t.base.CloseIdleConnections()
}

type clientBody struct {
	rc io.ReadCloser

	onCloseMu sync.Mutex
	onClose   func()
}

func (b *clientBody) Read(p []byte) (int, error) {
	var n int
	var err error = io.EOF
	if b.rc != nil {
		n, err = b.rc.Read(p)
	}
	if err == io.EOF {
		b.onCloseMu.Lock()
		defer b.onCloseMu.Unlock()
		if b.onClose != nil {
			// trigger only once
			onClose := b.onClose
			b.onClose = nil
			// defer here to match the Close method
			defer onClose()
		}
	}
	return n, err
}

func (b *clientBody) Close() error {
	b.onCloseMu.Lock()
	defer b.onCloseMu.Unlock()
	if b.onClose != nil {
		// trigger only once
		onClose := b.onClose
		b.onClose = nil
		// trigger after we've signalled the underlying io.Closer
		defer onClose()
	}
	if b.rc == nil {
		return nil
	}
	err := b.rc.Close()
	return err
}
