package chitin

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"reflect"
	"sync"
	"testing"

	"github.com/kr/logfmt"
)

func TestCallChain(t *testing.T) {
	// Each RPC is uniquely identified by the Trace-ID/Trace-Span pair
	// (barring uint64 collisions). All calls within an RPC tree share the
	// same Trace-ID, but will have different Trace-Spans.
	//
	// The Trace-Span identifies the RPC itself. It has the same value on the
	// caller as on the callee, but the event types emitted at those points is
	// different. When a service makes its 4th outbound call, it appends ".3"
	// (zero-indexed) to the Trace-Span it received from its caller. It would
	// send a Trace-Span of "XXX.3" if its caller specified a Trace-Span of
	// "XXX", and emit a trace event (tagged with the Trace-ID of course)
	// indicating that it started an outbound call for "XXX.3". When that call
	// is handled by the subservice, the subservice emits a trace event
	// indicating that it started processing an inbound call for "XXX.3". If
	// that subservice in turn makes any calls, the first will be labelled
	// "XXX.3.0".
	//
	// This test confirms that the Trace-Spans are propagated correctly.

	var wg sync.WaitGroup

	var buf bytes.Buffer
	lg := log.New(&buf, "", 0)

	servers := []*httptest.Server{
		httptest.NewUnstartedServer(&chainServer{
			wg: &wg, name: "A", callCount: 10, subcall: 5}), // sees one client request
		httptest.NewUnstartedServer(&chainServer{
			wg: &wg, name: "B", callCount: 15, subcall: 2}), // sees 10 calls, including ".5"
		httptest.NewUnstartedServer(&chainServer{
			wg: &wg, name: "C", callCount: 0, subcall: 0}), // sees 15 calls, including ".5.2"
	}
	for i := len(servers) - 1; i >= 0; i-- {
		srv := servers[i]
		handler := srv.Config.Handler
		srv.Config.Handler = Handler(handler, SetLogger(lg))
		srv.Listener = wrapListener(srv.Listener)

		srv.Start()
		defer srv.Close()

		if i == len(servers)-1 {
			continue
		}

		prev := servers[i+1]
		handler.(*chainServer).calleeBase = prev.URL
	}

	wg.Add(1)
	resp, err := http.Get(servers[0].URL + "/subcall")
	if err != nil {
		t.Fatalf("http request: %v", err)
	}
	resp.Body.Close()
	wg.Done()
	wg.Wait()

	var EndSpans []string
	sc := bufio.NewScanner(&buf)
	for sc.Scan() {
		type line struct {
			TraceID   string `logfmt:"trace-id"`
			TraceSpan string `logfmt:"trace-span"`
			Name      string `logfmt:"name"`
			Path      string `logfmt:"path"`
		}
		var l line
		err := logfmt.Unmarshal(sc.Bytes(), &l)
		if err != nil {
			t.Errorf("logfmt: %v", err)
			continue
		}
		if l.Name == "C" && l.Path == "/subcall" {
			EndSpans = append(EndSpans, l.TraceSpan)
		}
	}
	err = sc.Err()
	if err != nil {
		t.Errorf("scan: %v", err)
	}

	if have, want := EndSpans, []string{".5.2"}; !reflect.DeepEqual(have, want) {
		t.Errorf("EndSpans; %q != %q", have, want)
	}
}

type chainServer struct {
	wg         *sync.WaitGroup
	name       string
	callCount  int
	subcall    int
	calleeBase string
}

func (cs *chainServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	cs.wg.Add(1)
	defer cs.wg.Done()

	ctx := Context(w, r)
	printf(ctx, "name=%q path=%q", cs.name, r.URL.Path)

	if r.URL.Path != "/subcall" {
		return
	}

	client := Client(ctx)
	for i := 0; i < cs.callCount; i++ {
		suffix := "/"
		if i == cs.subcall {
			suffix = "/subcall"
		}
		resp, err := client.Get(cs.calleeBase + suffix)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		_, err = io.Copy(ioutil.Discard, resp.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		err = resp.Body.Close()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}
