package pd

import (
	"encoding/json"
	"fmt"
	"github.com/pingcap/schrodinger-test/once/drivers/pd/pd_ctl"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"testing"
)

var testBin string

func TestMain(m *testing.M) {
	var err error
	testBin, err = pd_ctl.InstallFromGo("master")
	if err != nil {
		fmt.Errorf("install pd-ctl error = %v", err)
		os.Exit(1)
	}
	os.Exit(m.Run())
}

func TestPD_ExistsStoreScheduler(t *testing.T) {
	type args struct {
		name    string
		storeID uint64
	}
	tests := []struct {
		name    string
		args    args
		want    bool
		wantErr bool
	}{
		{
			name: "exist-evict-leader-scheduler-1",
			args: args{
				name:    "evict-leader-scheduler",
				storeID: 1,
			},
			want:    true,
			wantErr: false,
		},
		{
			name: "not-exist-grant-leader-scheduler-1",
			args: args{
				name:    "grant-leader-scheduler",
				storeID: 1,
			},
			want:    false,
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schedulers := []string{
				"label-scheduler",
				"evict-leader-scheduler-1",
			}
			res, err := json.Marshal(schedulers)
			if err != nil {
				t.Errorf("Marshal scheduler error = %v, schedulers = %v", err, schedulers)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprintln(w, string(res))
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.ExistsStoreScheduler(tt.args.name, tt.args.storeID)
			if (err != nil) != tt.wantErr {
				t.Errorf("ExistsStoreScheduler() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("ExistsStoreScheduler() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_GetStoreLeaderMap(t *testing.T) {
	tests := []struct {
		name       string
		want       map[uint64][]Region
		wantErr    bool
		storesInfo StoresInfo
		regionsMap map[uint64][]Region
	}{
		{
			name: "get-store-region-map",
			want: map[uint64][]Region{
				1: []Region{
					{
						ID: 1,
						Leader: Peer{
							ID:      1,
							StoreID: 1,
						},
					},
				},
				2: []Region{
					{
						ID: 2,
						Leader: Peer{
							ID:      2,
							StoreID: 2,
						},
					},
					{
						ID: 3,
						Leader: Peer{
							ID:      3,
							StoreID: 2,
						},
					},
				},
			},
			wantErr: false,
			storesInfo: StoresInfo{
				Count: 2,
				Stores: []StoreInfo{
					{
						Store: Store{ID: 1},
					},
					{
						Store: Store{ID: 2},
					},
				},
			},
			regionsMap: map[uint64][]Region{
				1: []Region{
					{
						ID: 1,
						Leader: Peer{
							ID:      1,
							StoreID: 1,
						},
					},
					{
						ID: 2,
						Leader: Peer{
							ID:      2,
							StoreID: 2,
						},
					},
				},
				2: []Region{
					{
						ID: 2,
						Leader: Peer{
							ID:      2,
							StoreID: 2,
						},
					},
					{
						ID: 3,
						Leader: Peer{
							ID:      3,
							StoreID: 2,
						},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s1 := RegionsInfo{
				Count:   len(tt.regionsMap[1]),
				Regions: tt.regionsMap[1],
			}
			s2 := RegionsInfo{
				Count:   len(tt.regionsMap[2]),
				Regions: tt.regionsMap[2],
			}
			s1Data, err := json.Marshal(s1)
			if err != nil {
				t.Errorf("Marshal store 1 regions info error = %v, store 1 regions info: %v", err, s1)
			}
			s2Data, err := json.Marshal(s2)
			if err != nil {
				t.Errorf("Marshal store 2 regions info error = %v, store 2 regions info: %v", err, s2)
			}
			storesInfo, err := json.Marshal(tt.storesInfo)
			if err != nil {
				t.Errorf("Masrhal stores info error = %v, stores info: %v", err, storesInfo)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if r.URL.Path == "/pd/api/v1/stores" {
					fmt.Fprintln(w, string(storesInfo))
					return
				}
				if r.URL.Path == "/pd/api/v1/regions/store/1" {
					fmt.Fprintln(w, string(s1Data))
					return
				}
				if r.URL.Path == "/pd/api/v1/regions/store/2" {
					fmt.Fprintln(w, string(s2Data))
					return
				}
				fmt.Fprintln(w, "")
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.GetStoreRegionLeaderMap()
			if (err != nil) != tt.wantErr {
				t.Errorf("GetStoreRegionLeaderMap() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetStoreRegionLeaderMap() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_GetStoreRegionMap(t *testing.T) {
	tests := []struct {
		name       string
		want       map[uint64][]Region
		wantErr    bool
		storesInfo StoresInfo
	}{
		{
			name: "get-store-region-map",
			want: map[uint64][]Region{
				1: []Region{
					{
						ID: 1,
						Leader: Peer{
							ID:      1,
							StoreID: 1,
						},
					},
					{
						ID: 2,
						Leader: Peer{
							ID:      2,
							StoreID: 2,
						},
					},
				},
				2: []Region{
					{
						ID: 2,
						Leader: Peer{
							ID:      2,
							StoreID: 2,
						},
					},
					{
						ID: 3,
						Leader: Peer{
							ID:      3,
							StoreID: 2,
						},
					},
				},
			},
			wantErr: false,
			storesInfo: StoresInfo{
				Count: 2,
				Stores: []StoreInfo{
					{
						Store: Store{ID: 1},
					},
					{
						Store: Store{ID: 2},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s1 := RegionsInfo{
				Count:   len(tt.want[1]),
				Regions: tt.want[1],
			}
			s2 := RegionsInfo{
				Count:   len(tt.want[2]),
				Regions: tt.want[2],
			}
			s1Data, err := json.Marshal(s1)
			if err != nil {
				t.Errorf("Marshal store 1 regions info error = %v, store 1 regions info: %v", err, s1)
			}
			s2Data, err := json.Marshal(s2)
			if err != nil {
				t.Errorf("Marshal store 2 regions info error = %v, store 2 regions info: %v", err, s2)
			}
			storesInfo, err := json.Marshal(tt.storesInfo)
			if err != nil {
				t.Errorf("Masrhal stores info error = %v, stores info: %v", err, storesInfo)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if r.URL.Path == "/pd/api/v1/stores" {
					fmt.Fprintln(w, string(storesInfo))
					return
				}
				if r.URL.Path == "/pd/api/v1/regions/store/1" {
					fmt.Fprintln(w, string(s1Data))
					return
				}
				if r.URL.Path == "/pd/api/v1/regions/store/2" {
					fmt.Fprintln(w, string(s2Data))
					return
				}
				fmt.Fprintln(w, "")
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.GetStoreRegionMap()
			if (err != nil) != tt.wantErr {
				t.Errorf("GetStoreRegionMap() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetStoreRegionMap() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_ListRegions(t *testing.T) {
	tests := []struct {
		name    string
		want    []Region
		wantErr bool
	}{
		{
			name: "get-regions",
			want: []Region{
				{
					ID: 1,
					Leader: Peer{
						ID:      1,
						StoreID: 1,
					},
				},
				{
					ID: 2,
					Leader: Peer{
						ID:      2,
						StoreID: 1,
					},
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			regionsInfo := RegionsInfo{
				Count:   len(tt.want),
				Regions: tt.want,
			}
			res, err := json.Marshal(regionsInfo)
			if err != nil {
				t.Errorf("Marshal regions info error = %v, regions info: %v", err, regionsInfo)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprintln(w, string(res))
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.ListRegions()
			if (err != nil) != tt.wantErr {
				t.Errorf("ListRegions() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ListRegions() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_ListSchedulers(t *testing.T) {
	tests := []struct {
		name    string
		want    []string
		wantErr bool
	}{
		{
			name: "list-schedulers",
			want: []string{
				"label-scheduler",
				"evict-leader-scheduler-1",
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			res, err := json.Marshal(tt.want)
			if err != nil {
				t.Errorf("Marshal scheduler error = %v, schedulers = %v", err, tt.want)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprintln(w, string(res))
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.ListSchedulers()
			if (err != nil) != tt.wantErr {
				t.Errorf("ListSchedulers() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ListSchedulers() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_ListStores(t *testing.T) {
	tests := []struct {
		name    string
		want    []StoreInfo
		wantErr bool
	}{
		{
			name: "list-stores",
			want: []StoreInfo{
				{
					Store: Store{ID: 1},
				},
				{
					Store: Store{ID: 2},
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			storesInfo := StoresInfo{
				Count:  len(tt.want),
				Stores: tt.want,
			}
			res, err := json.Marshal(storesInfo)
			if err != nil {
				t.Errorf("Marshal stores info error = %v, stores info: %v", err, storesInfo)
			}
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprintln(w, string(res))
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			got, err := p.ListStores()
			if (err != nil) != tt.wantErr {
				t.Errorf("ListStores() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ListStores() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestPD_RemoveScheduler(t *testing.T) {
	type args struct {
		name string
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		{
			name: "remove-label-scheduler",
			args: args{
				name: "label-scheduler",
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprintln(w, "")
			}))
			defer ts.Close()
			pd_ctl.Setup(nil, testBin, ts.URL)
			p := &PD{}
			if err := p.RemoveScheduler(tt.args.name); (err != nil) != tt.wantErr {
				t.Errorf("RemoveScheduler() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}