package chitin

import (
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"code.justin.tv/common/chitin/internal/trace"
	"github.com/kr/logfmt"
)

func clientCancelHelper(t *testing.T, ctx context.Context, tickServer, tickClient, abort chan struct{}, run func()) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-tickServer:
		case <-abort:
		}
		select {
		case <-tickServer:
		case <-abort:
		}
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		select {
		case <-tickServer:
		case <-abort:
		}
		return
	}))
	defer srv.Close()
	defer close(abort)

	req, err := http.NewRequest("GET", srv.URL, nil)
	if err != nil {
		t.Fatalf("NewRequest; err = %q", err)
	}

	go func() {
		defer close(tickClient)

		resp, err := Client(ctx).Do(req)
		if err != nil {
			return
		}
		defer resp.Body.Close()
		tickClient <- struct{}{}

		_, err = io.Copy(ioutil.Discard, resp.Body)
		if err != nil {
			return
		}
		tickClient <- struct{}{}
		t.Errorf("body read succeeded unexpectedly")
	}()

	run()
}

func TestClientCancelHeaders(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)

	tickServer, tickClient := make(chan struct{}), make(chan struct{})
	abort := make(chan struct{})

	clientCancelHelper(t, ctx, tickServer, tickClient, abort, func() {
		// make sure the request make it to the server
		select {
		case tickServer <- struct{}{}:
		case <-ctx.Done():
			t.Errorf("ctx.Done triggered, but the request should still be active")
		}
		// cause the request to be cancelled
		cancel()
		// wait until client has read the header
		select {
		case <-tickClient:
		case <-time.After(1 * time.Second):
			t.Errorf("timeout triggered, but the request should have been canceled")
		}
		// wait until client has consumed the body
		select {
		case <-tickClient:
		case <-time.After(1 * time.Second):
			t.Errorf("timeout triggered, but the request should have been canceled")
		}
		// allow server to write header
		select {
		case tickServer <- struct{}{}:
		case <-time.After(1 * time.Second):
			t.Errorf("timeout triggered, but the request should have been canceled")
		}
	})
}

func TestClientCancelBody(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)

	tickServer, tickClient := make(chan struct{}), make(chan struct{})
	abort := make(chan struct{})

	clientCancelHelper(t, ctx, tickServer, tickClient, abort, func() {
		// make sure the request make it to the server
		select {
		case tickServer <- struct{}{}:
		case <-ctx.Done():
			t.Errorf("ctx.Done triggered, but the request should still be active")
		}
		// allow server to write header
		select {
		case tickServer <- struct{}{}:
		case <-ctx.Done():
			t.Errorf("ctx.Done triggered, but the request should still be active")
		}
		// wait until client has read the header
		select {
		case <-tickClient:
		case <-ctx.Done():
			t.Errorf("ctx.Done triggered, but the request should still be active")
		}
		// cause the request to be cancelled
		cancel()
		// wait until client has consumed the body
		select {
		case <-tickClient:
		case <-time.After(1 * time.Second):
			t.Errorf("timeout triggered, but the request should have been canceled")
		}
	})
}

func TestClientHeaders(t *testing.T) {
	ctx := context.Background()

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Pong", r.Header.Get("Ping"))
		w.Header().Set("Location", "mars")
		w.WriteHeader(http.StatusPaymentRequired)
	}))
	defer srv.Close()

	req, err := http.NewRequest("GET", srv.URL, nil)
	if err != nil {
		t.Fatalf("NewRequest; err = %q", err)
	}

	req.Header.Set("Ping", "FrankerZ")

	resp, err := Client(ctx).Do(req)
	if err != nil {
		t.Fatalf("Client.Do; err = %q", err)
	}
	defer resp.Body.Close()

	if have, want := resp.Header.Get("Pong"), "FrankerZ"; have != want {
		t.Errorf("Pong header; %q != %q", have, want)
	}
	if have, want := resp.Header.Get("Location"), "mars"; have != want {
		t.Errorf("Location header; %q != %q", have, want)
	}
	if have, want := resp.StatusCode, http.StatusPaymentRequired; have != want {
		t.Errorf("resp.StatusCode; %d != %d", have, want)
	}
}

func TestTraceSpanHeader(t *testing.T) {
	// generate a random trace id
	ctx := trace.ContextFromHeader(context.Background(), nil)

	// force the span id to be non-empty
	ctx0 := ctx
	for i := 0; i < 22+1; i++ {
		ctx = trace.NewSpan(ctx0)
	}
	ctx1 := ctx
	for i := 0; i < 33+1; i++ {
		ctx = trace.NewSpan(ctx1)
	}
	h := make(http.Header)
	trace.AugmentHeader(ctx, h)
	ctx = trace.ContextFromHeader(context.Background(), h)
	if have, want := trace.Logfmt(ctx), `".22.33"`; !strings.Contains(have, want) {
		t.Errorf("tx.SubID; %q does not contain %q", have, want)
		return
	}

	// our test requires a custom http client, arrange for it to be installed
	// for this test
	defer func(client *http.Client) {
		http.DefaultClient = client
	}(http.DefaultClient)
	http.DefaultClient = Client(ctx)

	type traceVals struct {
		TraceID   string `logfmt:"trace-id"`
		TraceSpan string `logfmt:"trace-span"`
	}

	var parent traceVals
	err := logfmt.Unmarshal([]byte(trace.Logfmt(ctx)), &parent)
	if err != nil {
		t.Errorf("logfmt unmarshal %v", err)
		return
	}

	children := make(chan traceVals, 5)

	for i := 0; i < cap(children); i++ {
		resp, err := testGet(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := Context(w, r)

			var child traceVals
			err := logfmt.Unmarshal([]byte(trace.Logfmt(ctx)), &child)
			if err != nil {
				t.Errorf("logfmt unmarshal %v", err)
				return
			}

			children <- child

			if have, want := child.TraceID, parent.TraceID; have != want {
				t.Errorf("transaction id was not successfully transmitted: %q != %q", have, want)
			}
			// if have, want := subTx.SubID, tx.SubID+suffix; have != want {
			// 	t.Errorf("transaction subspan was not successfully transmitted: %q != %q", have, want)
			// }
		}), "/")
		if err != nil {
			t.Errorf("http request failed: %v", err)
			return
		}
		resp.Body.Close()
	}
	close(children)

	spans := make(map[string]struct{}, cap(children))
	for child := range children {
		if _, ok := spans[child.TraceSpan]; ok {
			t.Errorf("duplicate span: %q", child.TraceSpan)
		}
		spans[child.TraceSpan] = struct{}{}
	}
}
