package chitin

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"runtime"
	"runtime/debug"
	"sync"
	"testing"
	"time"
)

func TestReservedHeader(t *testing.T) {
	for _, tt := range []struct {
		desc     string
		name     string
		reserved bool
	}{
		{name: "chitin", reserved: false, desc: "Headers must be normalized before calling reservedHeader"},
		{name: "Chitin", reserved: true, desc: "The prefix is reserved"},
		{name: "Chitin-Fu", reserved: true, desc: "Headers beginning with the prefix are reserved"},
		{name: "Chitintacular", reserved: false, desc: "The prefix check tokenizes on '-'"},
		{name: "Trace-Span", reserved: false, desc: "We don't reserve every header we use"},
	} {
		if have, want := reservedHeader(tt.name), tt.reserved; have != want {
			t.Errorf("reservedHeader(%q); %t != %t (%s)", tt.name, have, want, tt.desc)
		}
	}
}

func TestContextRetrieval(t *testing.T) {
	resp, err := testGet(t, func(
		w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "hi")

		ctx := Context(w, r)

		err := ctx.Err()
		if err != nil {
			t.Errorf("context is no longer active: %v", err)
		}

		arrived := getArrivalTime(ctx)
		now := time.Now()
		if delta := now.Sub(arrived); delta < 0 || delta > 1*time.Second {
			t.Errorf("improbable delta: %v", delta)
		}
	}, "/")

	if err != nil {
		t.Errorf("unexpected error: %v", err)
		return
	}
	defer resp.Body.Close()
}

func TestPanic(t *testing.T) {
	var buf bytes.Buffer
	log.SetOutput(&buf)
	defer func() {
		log.SetOutput(os.Stderr)
		if t.Failed() {
			t.Logf("logs:\n%s", buf.String())
		}
	}()

	resp, err := testGet(t, func(w http.ResponseWriter, r *http.Request) {
		panic("derp")
	}, "/")
	if err == nil {
		t.Errorf("panic in a handler should lead to an error on the client")
		resp.Body.Close()
		return
	}

	resp, err = testGet(t, func(w http.ResponseWriter, r *http.Request) {
		panic(nil)
	}, "/")
	if err == nil {
		t.Errorf("nil panic in a handler should lead to an error on the client")
		resp.Body.Close()
		return
	}

	resp, err = testGet(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		panic("derp")
	}, "/")
	if err != nil {
		t.Errorf("panic in a handler that has flushed its response need not "+
			"result in client error ... but it did: %v", err)
		return
	}
	resp.Body.Close()
}

func TestStatusCode(t *testing.T) {
	maker := func(status int) func(w http.ResponseWriter, r *http.Request) {
		return func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(status)
		}
	}

	mux := http.NewServeMux()
	mux.Handle("/", http.DefaultServeMux)
	mux.HandleFunc("/notfound", maker(http.StatusNotFound))
	mux.HandleFunc("/ok", maker(http.StatusOK))
	mux.HandleFunc("/error", maker(http.StatusInternalServerError))
	mux.HandleFunc("/teapot", maker(http.StatusTeapot))

	for _, tt := range []struct {
		url    string
		method string
		status int
	}{
		{url: "/notfound", status: http.StatusNotFound},
		{url: "/ok", status: http.StatusOK},
		{url: "/error", status: http.StatusInternalServerError},
		{url: "/teapot", status: http.StatusTeapot},

		{url: "/debug", status: http.StatusNotFound},
		{url: "/debug/", status: http.StatusNotFound},
		{url: "/debug/vars/", status: http.StatusNotFound},

		{url: "/debug/vars", status: http.StatusOK},
		{url: "/debug/pprof", status: http.StatusOK},
		{url: "/debug/pprof/", status: http.StatusOK},

		{url: "/ok", method: "GET", status: http.StatusOK},
		{url: "/ok", method: "POST", status: http.StatusOK},
		{url: "/ok", method: "PUT", status: http.StatusOK},
		{url: "/ok", method: "DELETE", status: http.StatusOK},
		{url: "/ok", method: "HEAD", status: http.StatusOK},
		{url: "/ok", method: "PATCH", status: http.StatusOK},
		{url: "/ok", method: "T", status: http.StatusMethodNotAllowed},
	} {
		if tt.method == "" {
			tt.method = "GET"
		}
		resp, err := testRequest(t, Handler(mux).ServeHTTP, tt.method, tt.url)
		if err != nil {
			t.Errorf("unexpected error: %v", err)
			continue
		}
		resp.Body.Close()
		if have, want := resp.StatusCode, tt.status; have != want {
			t.Errorf("testGet(%q) status = %d; want %d", tt.url, have, want)
		}
	}
}

func TestServerCloseBody(t *testing.T) {
	defer debug.SetMaxStack(debug.SetMaxStack(16 << 10))

	testGet(t, func(w http.ResponseWriter, r *http.Request) {
		r.Body.Close()
	}, "/")
}

func TestServerCloseNotifyLeak(t *testing.T) {
	const requests = 10

	done := make(chan struct{})

	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(
		w http.ResponseWriter, r *http.Request) {

		// trigger the problem case
		w.(http.CloseNotifier).CloseNotify()

		// allow the client to read full headers
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()

		// wait until we know we've got a bunch of concurrent requests
		<-done
	}))

	srv.Config.Handler = Handler(srv.Config.Handler)
	srv.Listener = wrapListener(srv.Listener)

	srv.Start()
	defer srv.Close()

	defer compareStacks(t, 10*time.Millisecond)()

	rt := &http.Transport{
		Dial: (&net.Dialer{
			Timeout: 1 * time.Second,
		}).Dial,
	}
	cl := &http.Client{
		Transport: rt,
	}

	var wg1, wg2 sync.WaitGroup
	for i := 0; i < requests; i++ {
		wg1.Add(1)
		wg2.Add(1)
		go func() {
			resp, err := cl.Get(srv.URL)
			if err != nil {
				wg1.Done()
				wg2.Done()
				return
			}
			defer resp.Body.Close()
			wg1.Done()
			io.Copy(ioutil.Discard, resp.Body)
			wg2.Done()
		}()
	}

	wg1.Wait()
	close(done)

	wg2.Wait()
	rt.CloseIdleConnections()
}

func testGet(t *testing.T, handlerFunc http.HandlerFunc, url string) (*http.Response, error) {
	return testRequest(t, handlerFunc, "GET", url)
}

func testRequest(t *testing.T, handlerFunc http.HandlerFunc, method string, url string) (*http.Response, error) {
	var wg sync.WaitGroup

	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(
		w http.ResponseWriter, r *http.Request) {
		wg.Add(1)
		defer wg.Done()
		handlerFunc(w, r)
	}))

	lg := log.New(ioutil.Discard, "", 0)
	srv.Config.Handler = Handler(srv.Config.Handler, SetLogger(lg))
	srv.Listener = wrapListener(srv.Listener)

	srv.Start()
	defer srv.Close()

	req, err := http.NewRequest(method, srv.URL+url, nil)
	if err != nil {
		return nil, err
	}

	wg.Add(1)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return resp, err
	}
	wg.Done()
	wg.Wait()
	return resp, nil
}

func BenchmarkStdlibServer(b *testing.B) {
	benchStdlibServer(b, func(w http.ResponseWriter, r *http.Request) {
	})
}

func BenchmarkStdlibServerHello(b *testing.B) {
	benchStdlibServer(b, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "hello %q", r.URL.String())
	})
}

func BenchmarkChitinServer(b *testing.B) {
	benchChitinServer(b, func(w http.ResponseWriter, r *http.Request) {
	})
}

func BenchmarkChitinServerHello(b *testing.B) {
	benchChitinServer(b, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "hello %q", r.URL.String())
	})
}

func BenchmarkChitinServerContextLookup(b *testing.B) {
	benchChitinServer(b, func(w http.ResponseWriter, r *http.Request) {
		Context(w, r)
	})
}

func benchStdlibServer(b *testing.B, fn http.HandlerFunc) {
	srv := httptest.NewUnstartedServer(fn)
	benchServer(b, srv)
}

func benchChitinServer(b *testing.B, fn http.HandlerFunc) {
	srv := httptest.NewUnstartedServer(fn)
	lg := log.New(ioutil.Discard, "", 0)

	srv.Config.Handler = Handler(srv.Config.Handler, SetLogger(lg))
	srv.Listener = wrapListener(srv.Listener)
	benchServer(b, srv)
}

func benchServer(b *testing.B, srv *httptest.Server) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stderr)

	srv.Start()
	defer srv.Close()

	const workers = 1 << 5

	def := http.DefaultTransport.(*http.Transport)
	tr := &http.Transport{
		Proxy:               def.Proxy,
		Dial:                def.Dial,
		TLSHandshakeTimeout: def.TLSHandshakeTimeout,
	}

	tr.MaxIdleConnsPerHost = workers
	client := &http.Client{}
	*client = *http.DefaultClient
	client.Transport = tr

	var wg sync.WaitGroup
	defer wg.Wait()
	work := make(chan struct{}, 1<<10)
	defer close(work)
	errs := make(chan error, 1)

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for _ = range work {
				resp, err := client.Get(srv.URL)
				if err != nil {
					select {
					case errs <- err:
					default:
					}
					return
				}
				io.Copy(ioutil.Discard, resp.Body)
				resp.Body.Close()
			}
		}()
	}

	for i := 0; i < b.N; i++ {
		select {
		case work <- struct{}{}:
		case err := <-errs:
			b.Errorf("request error: %v", err)
			return
		}
	}
}

func compareStacks(t *testing.T, pause time.Duration) func() {
	runtime.GC()
	time.Sleep(pause)

	before := runtime.NumGoroutine()
	beforeStacks := make([]byte, 64<<10)
	beforeStacks = beforeStacks[:runtime.Stack(beforeStacks, true)]

	return func() {
		runtime.GC()
		time.Sleep(pause)

		after := runtime.NumGoroutine()

		if before != after {
			t.Errorf("goroutine count changed during test (before=%d after=%d), suspect leak?",
				before, after)

			afterStacks := make([]byte, 256<<10)
			afterStacks = afterStacks[:runtime.Stack(afterStacks, true)]
			t.Logf("before stacks:\n%s\n", beforeStacks)
			t.Logf("after stacks:\n%s\n", afterStacks)
		}
	}
}
