package main

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

	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/prometheus/common/log"
)

type RegionAvailableDetector struct {
	db              *sql.DB
	database        string
	table           string
	maxResponseTime time.Duration
	sleepDur        time.Duration
	totalRows       int
	concurrency     int
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randString(maxLen int) string {
	n := rand.Intn(maxLen) + 1
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func NewRegionAvailableDetector(db *sql.DB, database, table string, totalRows, maxResponseSecs, sleepDur, concurrency int) (*RegionAvailableDetector, error) {
	return &RegionAvailableDetector{
		db:              db,
		database:        database,
		table:           table,
		totalRows:       totalRows,
		maxResponseTime: time.Duration(maxResponseSecs) * time.Second,
		sleepDur:        time.Duration(sleepDur) * time.Millisecond,
		concurrency:     concurrency,
	}, nil
}

func (d *RegionAvailableDetector) Initialize() error {
	// Create table
	if err := d.createTable(); err != nil {
		log.Error("create table fail")
		return err
	}

	// Load data
	if err := d.loadData(); err != nil {
		log.Error("load data fail")
		return err
	}

	return nil
}

func (d *RegionAvailableDetector) createTable() error {
	sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s.%s (id int(10) PRIMARY KEY, pad varchar(255))", d.database, d.table)
	_, err := d.db.Exec(sql)
	return err
}

func (d *RegionAvailableDetector) loadData() error {
	insertSQL := fmt.Sprintf("INSERT IGNORE INTO %s.%s VALUES", d.database, d.table)
	sql := insertSQL
	insertedRows := 0
	for i := 0; i < d.totalRows; i++ {
		sql += "(" + strconv.Itoa(i+1) + ", \"" + randString(255) + "\")"
		if i > 0 && i%100 == 0 {
			_, err := d.db.Exec(sql)
			if err != nil {
				return err
			}
			sql = insertSQL
			insertedRows = i + 1
		} else if i+1 < d.totalRows {
			sql += ","
		}
	}
	if insertedRows < d.totalRows {
		_, err := d.db.Exec(sql)
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *RegionAvailableDetector) Execute(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	var wg sync.WaitGroup
	for i := 0; i < d.concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				sql := fmt.Sprintf("UPDATE %s.%s SET pad=\"%s\" WHERE id = %d", d.database, d.table, randString(255), rand.Intn(d.totalRows)+1)
				start := time.Now()
				_, err := d.db.Exec(sql)
				if err != nil {
					client.Fatal("execute sql failed, ", sql, err)
				}
				end := time.Now()
				if end.Sub(start) > d.maxResponseTime {
					logStr := fmt.Sprintf("Query [%s] exceed max response time, takes %v sencods", sql, end.Sub(start).Seconds())
					client.Fatal(logStr)
				}

				select {
				case <-ctx.Done():
					return
				default:
					time.Sleep(d.sleepDur)
				}
			}

		}()
	}
	wg.Wait()
	cancel()
	d.db.Close()
}