package usher

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/pkg/errors"
	"github.com/twitchtv/twirp"
)

// usherRoundTripper is for use only with usherapi's client. This exists
// to deal with Usher's odd behavior. It supports the following options:
// - singlereply
// - extract
// See the usher service definition for a detailed description of each option.
func newUsherRoundTripper(base http.RoundTripper) (http.RoundTripper, error) {
	ann, err := readAnnotations(NewUsherServer(nil, nil).ServiceDescriptor())
	if err != nil {
		return nil, err
	}

	if base == nil {
		base = http.DefaultTransport
	}

	return &usherRoundTripper{
		base:              base,
		methodAnnotations: ann,
	}, nil
}

type usherRoundTripper struct {
	base              http.RoundTripper
	methodAnnotations map[string]*annotations
}

// RoundTrip round trips the request using the base RoundTripper and applies modifications to
// the request and response if the twirp method has option annotations set.
func (rt *usherRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	var err error

	methodName, ok := twirp.MethodName(req.Context())
	if !ok {
		return nil, errors.New("usherRoundTripper must only be used with a Twirp client")
	}

	if key := rt.methodAnnotations[methodName].extractRequestField; key != "" {
		req, err = extractRequestField(req, key)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to extract request field with key=%q", key)
		}
	}

	resp, err := rt.base.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	if rt.methodAnnotations[methodName].singlereply {
		resp, err = flattenSingleReply(resp)
		if err != nil {
			return nil, errors.Wrap(err, "failed to flatten single-reply response")
		}
	}

	return resp, nil
}

func flattenSingleReply(resp *http.Response) (*http.Response, error) {
	// only json content types are supported, just pass through
	if resp.Header.Get("Content-Type") != "application/json" {
		return resp, nil
	}

	respBody, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		return nil, errors.Wrap(err, "failed to read response body")
	}

	var val []json.RawMessage
	if err = json.Unmarshal(respBody, &val); err == nil && len(val) == 1 {
		// this is a json array, and its length is 1 so we unwrap it
		resp.Body = ioutil.NopCloser(bytes.NewReader(val[0]))
	} else {
		// all other cases pass through
		resp.Body = ioutil.NopCloser(bytes.NewReader(respBody))
	}

	return resp, nil
}

func extractRequestField(req *http.Request, key string) (*http.Request, error) {
	reqBody, err := ioutil.ReadAll(req.Body)
	_ = req.Body.Close()
	if err != nil {
		return nil, errors.Wrap(err, "failed to read request body")
	}

	var m map[string]json.RawMessage
	if err := json.Unmarshal(reqBody, &m); err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal request body")
	}

	val, found := m[key]
	if !found {
		return nil, errors.New("field not found in request body")
	}

	req.Body = ioutil.NopCloser(bytes.NewReader(val))
	req.ContentLength = int64(len(val))
	return req, nil
}
