// Copyright 2021 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package utils

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
)

func ReadAll(rows *sql.Rows) [][]string {
	defer func() {
		Try(rows.Close())
	}()

	data := make([][]string, 0)
	columns := Try(rows.Columns()).([]string)
	for rows.Next() {
		row := make([]interface{}, len(columns))
		scan := make([]interface{}, len(columns))
		for i := range row {
			scan[i] = &row[i]
		}
		Try(rows.Scan(scan...))

		rowstr := make([]string, len(columns))
		for i := range rowstr {
			switch r := row[i].(type) {
			case nil:
				rowstr[i] = "NULL"
			case []byte:
				rowstr[i] = string(r)
			default:
				rowstr[i] = fmt.Sprintf("%v", row[i])
			}
		}
		data = append(data, rowstr)
	}
	return data
}

func MustQuery(db *sql.DB, query string, args ...interface{}) *sql.Rows {
	return Try(db.Query(query, args...)).(*sql.Rows)
}

func WaitScheduled(db *sql.DB, table string, timeout time.Duration) error {
	return waitTableScheduleStateIn(db, table, []string{"SCHEDULED"}, timeout)
}

func WaitScheduledOrPending(db *sql.DB, table string, timeout time.Duration) error {
	return waitTableScheduleStateIn(db, table, []string{"SCHEDULED", "PENDING"}, timeout)
}

func waitTableScheduleStateIn(db *sql.DB, table string, states []string, timeout time.Duration) error {
	startTime := time.Now()
	statesMap := make(map[string]bool, len(states))
	for _, state := range states {
		statesMap[strings.ToUpper(state)] = true
	}

	for {
		rows, err := db.Query(fmt.Sprintf("show placement for table %s", table))
		if err == nil {
			done := true
			var target, placement, state string
			for rows.Next() {
				Try(rows.Scan(&target, &placement, &state))
				if ok, _ := statesMap[strings.ToUpper(state)]; !ok {
					done = false
				}
			}

			if !done {
				log.Info("Still waiting for schedule state",
					zap.String("target", target),
					zap.String("placement", placement),
					zap.String("state", state),
					zap.String("expect", strings.Join(states, ",")),
				)
			}

			if !done && time.Now().Sub(startTime) > timeout {
				return errors.New("wait timeout")
			}

			if done {
				break
			}
		} else {
			log.Info("PD not available", zap.Error(err))
		}

		time.Sleep(time.Second)
	}
	return nil
}

type TiKVPeer struct {
	ID      uint64
	StoreID uint64
	Role    string
	Labels  map[string]string
}

func (s *TiKVPeer) Match(constraint interface{}) bool {
	if orConstraints, ok := constraint.([]string); ok {
		for _, c := range orConstraints {
			if s.Match(c) {
				return true
			}
		}
		return false
	}

	constraint = strings.TrimSpace(constraint.(string))
	if constraint == "" {
		return true
	}

	items := strings.Split(constraint.(string), ",")
	for _, item := range items {
		item = strings.TrimSpace(item)
		pair := strings.Split(item, "=")
		key := strings.TrimSpace(pair[0])
		exclude := key[0] == '-'
		key = key[1:]
		value := strings.TrimSpace(pair[1])
		v, ok := s.Labels[key]
		if exclude && ok {
			return false
		}

		if !exclude && (!ok || v != value) {
			return false
		}
	}
	return true
}

func GetAllStoreLabels(db *sql.DB) map[uint64]map[string]string {
	rows := MustQuery(db, "select STORE_ID, LABEL from INFORMATION_SCHEMA.TIKV_STORE_STATUS")
	stores := make(map[uint64]map[string]string, 0)
	for rows.Next() {
		var storeID uint64
		var labels string
		Try(rows.Scan(&storeID, &labels))
		labelsList := make([]map[string]string, 0)
		Try(json.Unmarshal([]byte(labels), &labelsList))
		labelsMap := make(map[string]string)
		for _, label := range labelsList {
			labelsMap[label["key"]] = label["value"]
		}
		stores[storeID] = labelsMap
	}
	return stores
}

type Regions map[uint64][]*TiKVPeer

func (r Regions) CheckRolePlacement(role, constraint interface{}, cnt int) Regions {
	regions := make(map[uint64][]*TiKVPeer)
	for regionID, peers := range r {
		leftPeers := make([]*TiKVPeer, 0)
		matchCnt := 0
		for _, peer := range peers {
			if peer.Match(constraint) && matchCnt < cnt && peer.Role == role {
				matchCnt++
			} else if cnt > 0 {
				leftPeers = append(leftPeers, peer)
			}
		}
		regions[regionID] = leftPeers
		if matchCnt != cnt {
			panic(errors.New(fmt.Sprintf(
				"CheckRolePlacement failed, expect %s has %d peers in region '%d' with constraints '%s', actual: %d",
				role, cnt, regionID, constraint, matchCnt,
			)))
		}
	}
	return regions
}

func (r Regions) CheckLeaderPlacement(constraint interface{}) Regions {
	return r.CheckRolePlacement("leader", constraint, 1)
}

func (r Regions) CheckFollowerPlacement(constraint interface{}, cnt int) Regions {
	return r.CheckRolePlacement("follower", constraint, cnt)
}

func (r Regions) CheckLearnerPlacement(constraint interface{}, cnt int) Regions {
	return r.CheckRolePlacement("learner", constraint, cnt)
}

func (r Regions) CheckFollowerRegions(cnt int, regions ...string) Regions {
	constraints := make([]string, len(regions))
	for _, region := range constraints {
		constraints = append(constraints, fmt.Sprintf("+region=%s,-engine=tiflash", region))
	}
	return r.CheckFollowerPlacement(constraints, cnt)
}

func (r Regions) CheckNoPeersLeft() {
	for regionID, peers := range r {
		if len(peers) != 0 {
			panic(errors.New(fmt.Sprintf("expect no peers left, but %d in region %d", len(peers), regionID)))
		}
	}
}

func CheckRegions(ctx *SuiteContext, table, partition string, fn func(r Regions)) {
	maxRetry := 10
	var err error
	for i := 0; i <= maxRetry; i++ {
		regions := GetRegions(ctx, table, partition)
		func() {
			defer func() {
				err = nil
				if r := recover(); r != nil {
					if e, ok := r.(error); ok {
						err = e
					} else {
						panic(r)
					}
				}
			}()
			fn(regions)
		}()

		if err == nil || i == maxRetry {
			break
		}

		log.Info("CheckRegions failed, maybe not stable yet retry",
			zap.String("error", err.Error()),
			zap.Int("retry", i+1),
			zap.Any("regions", regions),
		)
		time.Sleep(time.Second * time.Duration(2*i+2))
	}

	if err != nil {
		log.Error("CheckRegions failed", zap.Error(err))
	}
	Expect(err).Should(BeNil())
}

func GetRegions(ctx *SuiteContext, table, partition string) Regions {
	db := ctx.DB
	pd := ctx.PD
	dbName := ReadAll(MustQuery(db, "select database()"))[0][0]
	items := strings.SplitN(table, ".", 2)
	if len(items) == 2 {
		dbName = items[0]
		table = items[1]
	}

	regions := make(map[uint64][]*TiKVPeer)

	var tableIDStr string
	if partition == "" {
		tableIDStr = ReadAll(MustQuery(db,
			"select TIDB_TABLE_ID from INFORMATION_SCHEMA.TABLES "+
				"where TABLE_SCHEMA=? and TABLE_NAME=?", dbName, table))[0][0]
	} else {
		tableIDStr = ReadAll(MustQuery(db,
			"select TIDB_PARTITION_ID from INFORMATION_SCHEMA.PARTITIONS "+
				"where TABLE_SCHEMA=? and TABLE_NAME=? and PARTITION_NAME=?", dbName, table, partition))[0][0]
	}

	tableID := Try(strconv.ParseUint(tableIDStr, 10, 64)).(uint64)
	storeLabels := GetAllStoreLabels(db)
	rows := ReadAll(MustQuery(db, fmt.Sprintf("show table %s.%s regions", dbName, table)))
	for _, row := range rows {
		regionID := Try(strconv.ParseUint(row[0], 10, 64)).(uint64)
		startKey := row[1]
		endKey := row[2]
		if endKey != "" && endKey <= fmt.Sprintf("t_%d_", tableID) {
			continue
		}

		if startKey >= fmt.Sprintf("t_%d_", tableID+1) {
			continue
		}

		result := make(map[string]interface{})
		bs := Try(pd.GetJsonBytes(fmt.Sprintf("/pd/api/v1/region/id/%d", regionID))).([]byte)
		Try(json.Unmarshal(bs, &result))
		leaderID := uint64(result["leader"].(map[string]interface{})["id"].(float64))

		peers := make([]*TiKVPeer, 0)
		for _, p := range result["peers"].([]interface{}) {
			pInfo := p.(map[string]interface{})
			peerID := uint64(pInfo["id"].(float64))
			storeID := uint64(pInfo["store_id"].(float64))
			roleName := pInfo["role_name"].(string)

			role := "follower"
			if peerID == leaderID {
				role = "leader"
			} else if roleName != "Voter" {
				role = "learner"
			}

			peers = append(peers, &TiKVPeer{
				ID:      peerID,
				StoreID: storeID,
				Role:    role,
				Labels:  storeLabels[storeID],
			})
		}

		regions[regionID] = peers
	}

	Expect(len(regions) > 0).Should(BeTrue(), fmt.Sprintf("table %s, partition %s has no regions", table, partition))
	return regions
}

func GetDBWithUser(tc resource.Addressable, dbName, user string) (*sql.DB, error) {
	addr, err := tc.ServiceURL(resource.DBAddr)
	if err != nil {
		return nil, err
	}
	addr.User = url.User(user)

	db, err := sql.Open("mysql", resource.BuildDSN(addr, dbName, "timeout=2m"))
	if err != nil {
		return nil, err
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	return db, nil
}

func Sort(rows [][]string) [][]string {
	sort.Slice(rows, func(i, j int) bool {
		a := rows[i]
		b := rows[j]
		for i := range a {
			if a[i] < b[i] {
				return true
			} else if a[i] > b[i] {
				return false
			}
		}
		return false
	})
	return rows
}