package staleread

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

	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
)

// Config for stale read test
type Config struct {
	DBName          string
	TableName       string
	TotalRows       int
	Concurrency     int
	RequestInterval time.Duration
	MaxStaleness    int
}

type staleRead struct {
	*Config
	db      *sql.DB
	counter int64
}

func RandStringBytes(bytes []byte, rnd *rand.Rand) {
	letterBytes := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	for i := range bytes {
		bytes[i] = letterBytes[rnd.Intn(len(letterBytes))]
	}
}
func (s *staleRead) Start(ctx context.Context, tc resource.TiDBCluster, cfg interface{}) error {
	log.Info("[stale read] start to test...")
	var wg sync.WaitGroup
	// Wait `MaxStaleness` elapsed so all read can reture data
	time.Sleep(time.Duration(s.MaxStaleness) * time.Second)
	// Read data
	for i := 0; i < s.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			Try(s.CompareStaleStrongRead(ctx, tc))
		}()
	}
	wg.Wait()
	return nil
}

func (s *staleRead) CompareStaleStrongRead(ctx context.Context, tc resource.TiDBCluster) error {
	conn1, err := s.db.Conn(ctx)
	if err != nil {
		return err
	}
	conn2, err := s.db.Conn(ctx)
	if err != nil {
		return err
	}
	count := 0
	for {
		// prepare id
		ids := make([]int, rand.Intn(99)+1)
		for i := range ids {
			ids[i] = rand.Intn(s.TotalRows) + 1
		}
		ago := rand.Intn(s.MaxStaleness)
		now := time.Now().UTC()
		timeAgo := now.Add(time.Duration(-ago) * time.Second)
		stale, ts, err1 := s.staleRead(ctx, conn1, timeAgo, ids)
		strong, err2 := s.strongRead(ctx, conn2, ts, ids)
		if err1 == nil && err2 == nil {
			// The result should be the same
			if len(stale) != len(strong) {
				sql := fmt.Sprintf("select * from %s.%s where id in (%s)", s.DBName, s.Config.TableName, strings.Trim(strings.Join(strings.Fields(fmt.Sprint(ids)), ", "), "[]"))
				log.Fatal(fmt.Sprintf("got different number of resluts at %d ago, ids: %d, stale: %d, strong: %d, stale res: %v, strong res: %v, sql: %s", ago, len(ids), len(stale), len(strong), stale, strong, sql))
				return fmt.Errorf("got different number of resluts at %d ago, ids: %d, stale: %d, strong: %d, stale res: %v, strong res: %v", ago, len(ids), len(stale), len(strong), stale, strong)
			}
			for id, pad := range strong {
				if stalePad, ok := stale[id]; !ok || stalePad != pad {
					log.Fatal(fmt.Sprintf("got different resluts for id %d at %ds ago, tso: %v, stale read: %s, strong read: %s", id, ago, ts, stale[id], pad))
					return fmt.Errorf("got different resluts for id %d at %ds ago, tso: %v, stale read: %s, strong read: %s", id, ago, ts, stale[id], pad)
				}
			}
			if count%10 == 0 {
				sql := fmt.Sprintf("select * from %s.%s where id in (%s)", s.DBName, s.Config.TableName, strings.Trim(strings.Join(strings.Fields(fmt.Sprint(ids)), ", "), "[]"))
				log.Info(fmt.Sprintf("Stale read %d seconds ago equal to strong read, sql: %s", ago, sql))
			}
		} else {
			log.Warn(fmt.Sprintf("read return error, stale: %v, strong: %v", err1, err2))
		}

		select {
		case <-ctx.Done():
			return nil
		default:
			time.Sleep(s.RequestInterval)
			count++
		}
	}
}

func (s *staleRead) staleRead(ctx context.Context, conn *sql.Conn, t time.Time, ids []int) (map[int]string, uint64, error) {
	setSql := fmt.Sprintf("SET TRANSACTION READ ONLY AS OF TIMESTAMP '%s'", t.Format("2006-01-02 15:04:05"))
	log.Info(setSql)
	if _, err := conn.ExecContext(ctx, setSql); err != nil {
		log.Warn(fmt.Sprintf("failed to set stale read-only transaction, time: %s, err: %v", t.Format("2006-01-02 15:04:05"), err))
		return nil, 0, err
	}

	txn, err := conn.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		log.Warn(fmt.Sprintf("failed to start a read-only transaction, err: %v", err))
		return nil, 0, err
	}
	defer txn.Rollback()

	var tso uint64
	if err = txn.QueryRow("select @@tidb_current_ts").Scan(&tso); err != nil {
		log.Warn(fmt.Sprintf("failed to get tidb_current_ts, err: %v", err))
		return nil, 0, err
	}

	idsStr := strings.Trim(strings.Join(strings.Fields(fmt.Sprint(ids)), ", "), "[]")
	rows, err := txn.QueryContext(ctx, fmt.Sprintf("select * from %s.%s where id in (%s)", s.DBName, s.Config.TableName, idsStr))
	if err != nil {
		return nil, tso, err
	}
	defer rows.Close()
	results := make(map[int]string, len(ids))
	for rows.Next() {
		var id int
		var k int
		var c string
		var pad string
		if err = rows.Scan(&id, &k, &c, &pad); err != nil {
			return nil, tso, err
		}
		results[id] = pad
	}

	if err := txn.Commit(); err != nil {
		log.Warn(fmt.Sprintf("failed to commit a read-only transaction, err: %v", err))
		return nil, tso, err
	}
	return results, tso, nil
}

func (s *staleRead) strongRead(ctx context.Context, conn *sql.Conn, tso uint64, ids []int) (map[int]string, error) {
	if tso == 0 {
		return nil, errors.New("tso can't be zero")
	}
	if _, err := conn.ExecContext(ctx, fmt.Sprintf("set @@tidb_snapshot = '%d'", tso)); err != nil {
		log.Warn(fmt.Sprintf("failed to set tidb_snapshot, tso: %d, err: %v", tso, err))
		return nil, err
	}

	idsStr := strings.Trim(strings.Join(strings.Fields(fmt.Sprint(ids)), ", "), "[]")
	rows, err := conn.QueryContext(ctx, fmt.Sprintf("select * from %s.%s where id in (%s)", s.DBName, s.Config.TableName, idsStr))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	results := make(map[int]string, len(ids))
	for rows.Next() {
		var id int
		var k int
		var c string
		var pad string
		if err = rows.Scan(&id, &k, &c, &pad); err != nil {
			return nil, err
		}
		results[id] = pad
	}
	return results, nil
}