package qhad

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"math/rand"
	"time"

	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/workloadnode/cmdstr"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/endless/testcase/oltp/scenario/params"
)

func InitializeFullScan001Params(ctx context.Context, tidbClient *tidb.Client) (*params.FullScan, error) {
	params := params.FullScan{
		RunTime:    1 * time.Hour,
		WarmUpTime: 7 * time.Minute,
	}
	return &params, nil
}

func FullScanPartitionsRepeatedly(
	ctx context.Context,
	tc resource.TiDBCluster,
	opts cmdstr.QHAdOptions,
) error {

	dbName := opts.DBName

	dbUrl, err := tc.ServiceURL(resource.DBAddr)
	if err != nil {
		return err
	}

	db, err := sql.Open("mysql", resource.BuildDSN(dbUrl, dbName, "timeout=300m", "readTimeout=300m", "writeTimeout=300m"))
	if err != nil {
		return err
	}
	defer db.Close()

	tables := oltppkg.GetAllTablesNames(ctx, tc, dbName)
	tableNum := len(tables)
	log.Info("start full scan for partition tables Repeatedly")
	var round int
	for {
		select {
		case <-ctx.Done():
			log.Info("full partition scan ctx done.")
			return nil
		default:
		}
		for i := 0; i < tableNum; i++ {
			tableName := tables[i]
			partitionInfo, err := GetPartitionInfo(ctx, db, dbName, tableName)
			if err != nil {
				if !errors.Is(err, context.DeadlineExceeded) {
					return err
				} else {
					log.Info("GetPartitionInfo ctx done.")
					return nil
				}
			}
			//scan all partition table
			sql1 := fmt.Sprintf("select count(*) from %s;", tableName)

			//get a partition random
			partitions := partitionInfo["pName"]
			var randI int
			if len(partitions) <= 1 {
				// always have this partition
				randI = 0
			} else {
				// do not use the oldest partition because it may be droped.
				randI = rand.Intn(len(partitions)-1) + 1
			}
			pName := partitionInfo["pName"][randI]
			sql2 := fmt.Sprintf("select count(*) from %s partition (%s)", tableName, pName)

			sqls := []string{sql1, sql2}
			round++
			err = compareDynamicWithStatic(ctx, db, sqls, round)
			if err != nil {
				if !errors.Is(err, context.DeadlineExceeded) {
					return err
				}
			}
		}
	}
}

func GetPartitionInfo(ctx context.Context, db *sql.DB, dbName string, tableName string) (map[string][]string, error) {
	rows, err := db.QueryContext(ctx, "select partition_name,PARTITION_DESCRIPTION from information_schema.partitions where table_schema = ? and table_name = ? order by partition_name asc;", dbName, tableName)
	if err != nil {
		return nil, err
	}
	var pName string
	var pDescription string
	var partitionInfo = make(map[string][]string, 10)

	for rows.Next() {
		if err := rows.Scan(&pName, &pDescription); err != nil {
			return nil, err
		}
		partitionInfo["pName"] = append(partitionInfo["pName"], pName)
		partitionInfo["pDescription"] = append(partitionInfo["pDescription"], pDescription)
	}
	log.Info("partitionInfo ", zap.Strings("pName", partitionInfo["pName"]), zap.Strings("pDescription", partitionInfo["pDescription"]))
	return partitionInfo, nil
}

func compareDynamicWithStatic(ctx context.Context, db *sql.DB, sqls []string, round int) error {
	log.Info("[checker] check result of dynamic and static mode ", zap.Int("ROUND", round))
	for _, sql := range sqls {
		var c1, c2 int
		txn, _ := db.BeginTx(ctx, nil)
		log.Info("[workload] ", zap.String("sql", sql))
		if _, err := txn.ExecContext(ctx, "set tidb_partition_prune_mode='dynamic';"); err != nil {
			return err
		}
		if err := txn.QueryRow(sql).Scan(&c1); err != nil {
			return err
		}
		log.Info(fmt.Sprintf("c1: %d", c1))

		if _, err := txn.ExecContext(ctx, "set tidb_partition_prune_mode='static';"); err != nil {
			return err
		}
		if err := txn.QueryRow(sql).Scan(&c2); err != nil {
			return err
		}
		log.Info(fmt.Sprintf("c2: %d", c2))
		if c1 != c2 {
			return fmt.Errorf("c1 should equal to c2")
		}

		txn.Rollback()
	}
	return nil
}