package twirpmetric

import (
	"bytes"
	"context"
	"io/ioutil"
	"net/http"

	"code.justin.tv/video/metrics-middleware/v2/operation"
)

// Client wraps a http.RoundTripper with an operation.Starter.
type Client struct {
	Transport http.RoundTripper
	Starter   *operation.Starter
}

// RoundTrip executes a single HTTP transaction wrapped in a client operation.
func (client *Client) RoundTrip(request *http.Request) (*http.Response, error) {
	opName := getOp(request.Context())
	opName.Kind = operation.KindClient
	ctx, op := client.Starter.StartOp(request.Context(), opName)

	response, err := client.Transport.RoundTrip(request.WithContext(ctx))

	if err != nil {
		handleRoundTripFailure(ctx, op, err)
	} else if response.StatusCode == 200 { // Generated twirp clients accept only 200 as a success.
		handleRoundTripSuccess(ctx, op, response)
	} else {
		handleRoundTripErrorResponse(ctx, op, response)
	}

	return response, err
}

// handleRoundTripFailure ends the operation immediately since the response body is unusable.
func handleRoundTripFailure(ctx context.Context, op *operation.Op, err error) {
	// It failed to actually do the request.  Generated twirp clients report this as an internal error.
	setStatusWithContext(ctx, op, operation.Status{
		Code:    13, // Internal
		Message: "failed to do request: " + err.Error(),
	})
	op.End()
}

// handleRoundTripSuccess adds a wrapper to the response body and ends the operation when the response body read is complete.
func handleRoundTripSuccess(ctx context.Context, op *operation.Op, response *http.Response) {
	onReadComplete := func(responseBodyBytes []byte, readErr error) {
		status := operation.Status{}
		if readErr != nil && readErr != errIncompleteRead {
			// There was an error reading the body.  Generated twirp clients report this as an internal error.
			// An incomplete read is treated as a success, assuming that the caller did not need the rest of the response.
			status = operation.Status{
				Code:    13, // Internal
				Message: "failed to read response body: " + readErr.Error(),
			}
		}
		setStatusWithContext(ctx, op, status)
		op.End()
	}

	monitor := &readMonitor{readCloser: response.Body, captureBytesRead: false, onReadComplete: onReadComplete}
	response.Body = monitor
}

// handleRoundTripErrorResponse adds a wrapper to the response body and ends the operation when the response body read is complete.
// It collects a copy of the body to get the twirp error code.
func handleRoundTripErrorResponse(ctx context.Context, op *operation.Op, response *http.Response) {
	onReadComplete := func(responseBodyBytes []byte, readErr error) {
		defer op.End()

		if readErr == nil {
			errorResponse := &http.Response{
				StatusCode: response.StatusCode,
				Header:     response.Header,
				Body:       ioutil.NopCloser(bytes.NewReader(responseBodyBytes)),
			}
			twirpError := errorFromResponse(errorResponse)

			setStatusWithContext(ctx, op, statusFromTwirp(twirpError.Code()))
		} else {
			// There was an error reading the body.  Generated twirp clients report this as an internal error.
			// If the error response is not completely read, an internal error is reported since twirp clients
			// aren't expected to do that.
			setStatusWithContext(ctx, op, operation.Status{
				Code:    13, // Internal
				Message: "failed to read response body: " + readErr.Error(),
			})
		}
	}

	monitor := &readMonitor{readCloser: response.Body, captureBytesRead: true, onReadComplete: onReadComplete}
	response.Body = monitor
}

// setStatusWithContext sets the op status to Canceled or DeadlineExceeded if the context is one of those,
// otherwise it is set to the provided status.
func setStatusWithContext(ctx context.Context, op *operation.Op, status operation.Status) {
	switch err := ctx.Err(); err {
	case context.Canceled:
		op.SetStatus(operation.Status{
			Code:    1, // Canceled
			Message: err.Error(),
		})
	case context.DeadlineExceeded:
		op.SetStatus(operation.Status{
			Code:    4, // DeadlineExceeded
			Message: err.Error(),
		})
	default:
		op.SetStatus(status)
	}
}
