import (
	"encoding/gob"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	redistore "gopkg.in/boj/redistore.v1"

	"github.com/antonlindstrom/pgstore"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/sessions"
)

func benchGorilla(b *testing.B, store sessions.Store) {
	setupHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s, err := store.Get(r, "session")
		if err != nil {
			b.Fatal(err)
		}
		s.Values["counter"] = 1
		err = s.Save(r, w)
		if err != nil {
			b.Fatal(err)
		}
	})

	benchHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s, err := store.Get(r, "session")
		if err != nil {
			b.Fatal(err)
		}
		i, ok := s.Values["counter"].(int)
		if ok == false {
			b.Fatal(err)
		}
		s.Values["counter"] = i + 1
		err = s.Save(r, w)
		if err != nil {
			b.Fatal(err)
		}
	})

	w := httptest.NewRecorder()
	r := new(http.Request)

	setupHandler.ServeHTTP(w, r)
	sessionCookie := w.Header().Get("Set-Cookie")

	r = new(http.Request)
	r.Header = make(http.Header)
	r.Header.Add("Cookie", sessionCookie)

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		var nr http.Request
		nr = *r
		benchHandler.ServeHTTP(w, &nr)
	}
}

type User struct {
	Name string
	Age  int
}

func benchGorillaObject(b *testing.B, store sessions.Store) {
	gob.Register(new(User))

	setupHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s, err := store.Get(r, "session")
		if err != nil {
			b.Fatal(err)
		}
		s.Values["user"] = &User{"Alex", 33}
		err = s.Save(r, w)
		if err != nil {
			b.Fatal(err)
		}
	})

	benchHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s, err := store.Get(r, "session")
		if err != nil {
			b.Fatal(err)
		}
		user, ok := s.Values["user"].(*User)
		if ok == false {
			b.Fatal(err)
		}
		user.Age = user.Age + 1
		err = s.Save(r, w)
		if err != nil {
			b.Fatal(err)
		}
	})

	w := httptest.NewRecorder()
	r := new(http.Request)

	setupHandler.ServeHTTP(w, r)
	sessionCookie := w.Header().Get("Set-Cookie")

	r = new(http.Request)
	r.Header = make(http.Header)
	r.Header.Add("Cookie", sessionCookie)

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		var nr http.Request
		nr = *r
		benchHandler.ServeHTTP(w, &nr)
	}
}

func BenchmarkGorillaCookies(b *testing.B) {
	store := sessions.NewCookieStore(
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
		[]byte("911182cec2f206986c8c82440adb7d17"),
	)
	benchGorilla(b, store)
}

func BenchmarkGorillaRedis(b *testing.B) {
	pool := redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", os.Getenv("SESSION_REDIS_TEST_ADDR"))
		if err != nil {
			return nil, err
		}
		return conn, err
	}, 50)
	defer pool.Close()

	store, err := redistore.NewRediStoreWithPool(
		pool,
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
	)
	if err != nil {
		b.Fatal(err)
	}

	benchGorilla(b, store)
}

func BenchmarkGorillaPostgres(b *testing.B) {
	store, err := pgstore.NewPGStore(
		os.Getenv("SESSION_PG_TEST_DSN"),
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
	)
	if err != nil {
		b.Fatal(err)
	}
	defer store.Close()
	defer store.StopCleanup(store.Cleanup(time.Minute * 5))

	benchGorilla(b, store)
}

func BenchmarkGorillaObjectCookies(b *testing.B) {
	store := sessions.NewCookieStore(
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
		[]byte("911182cec2f206986c8c82440adb7d17"),
	)
	benchGorillaObject(b, store)
}

func BenchmarkGorillaObjectRedis(b *testing.B) {
	pool := redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", os.Getenv("SESSION_REDIS_TEST_ADDR"))
		if err != nil {
			return nil, err
		}
		return conn, err
	}, 50)
	defer pool.Close()

	store, err := redistore.NewRediStoreWithPool(
		pool,
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
	)
	if err != nil {
		b.Fatal(err)
	}

	benchGorillaObject(b, store)
}

func BenchmarkGorillaObjectPostgres(b *testing.B) {
	store, err := pgstore.NewPGStore(
		os.Getenv("SESSION_PG_TEST_DSN"),
		[]byte("f71dc7e58abab014ddad2652475056f185164d262869c8931b239de52711ba87"),
	)
	if err != nil {
		b.Fatal(err)
	}
	defer store.Close()
	defer store.StopCleanup(store.Cleanup(time.Minute * 5))

	benchGorillaObject(b, store)
}