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 ConflictsMaker struct {
	db                   *sql.DB
	database             string
	table                string
	readMaxResponseTime  time.Duration
	writeMaxResponseTime time.Duration
	writeSleepDur        time.Duration
	readSleepDur         time.Duration
	totalRows            int
	readConcurrency      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 NewConflictsMaker(db *sql.DB, database, table string, totalRows, writeMaxResponseSecs,
	readMaxResponseSecs, writeSleepDur, readSleepDur, readConcurrency int) (*ConflictsMaker, error) {
	return &ConflictsMaker{
		db:                   db,
		database:             database,
		table:                table,
		totalRows:            totalRows,
		writeMaxResponseTime: time.Duration(writeMaxResponseSecs) * time.Second,
		readMaxResponseTime:  time.Duration(readMaxResponseSecs) * time.Second,
		writeSleepDur:        time.Duration(writeSleepDur) * time.Millisecond,
		readSleepDur:         time.Duration(readSleepDur) * time.Millisecond,
		readConcurrency:      readConcurrency,
	}, nil
}

func (d *ConflictsMaker) 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 *ConflictsMaker) createTable() error {
	sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s.%s (id int(10) PRIMARY KEY, pad varchar(255), num int(10), "+
		"update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, "+
		"KEY idx_utime(update_time), KEY idx_pad(pad))", d.database, d.table)
	_, err := d.db.Exec(sql)
	return err
}

func (d *ConflictsMaker) 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(64) + "\"," + strconv.Itoa(rand.Intn(d.totalRows)) + ", now())"
		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 *ConflictsMaker) Execute(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	var wg sync.WaitGroup

	// Big transaction which will update all rows
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			sql := fmt.Sprintf("UPDATE %s.%s SET num=%d", d.database, d.table, 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.writeMaxResponseTime {
				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.writeSleepDur)
			}
		}
	}()

	// select a row
	for i := 0; i < d.readConcurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				sql := fmt.Sprintf("SELECT * FROM %s.%s WHERE id = %d", d.database, d.table, 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.readMaxResponseTime {
					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.readSleepDur)
				}
			}
		}()
	}
	wg.Wait()
	cancel()
	d.db.Close()
}