// Copyright 2018 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package main

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

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
)

type ReverseScanCase struct {
	totalRows   int
	concurrency int
}

type Config struct {
	TotalRows   int `toml:total_rows`
	Concurrency int `toml:concurrency`
}

func NewReverseScanCase(cfg *Config) *ReverseScanCase {
	b := &ReverseScanCase{
		totalRows:   cfg.TotalRows,
		concurrency: cfg.Concurrency,
	}
	return b
}

func (c *ReverseScanCase) String() string {
	return "reverse_scan"
}

func (c *ReverseScanCase) Initialize(ctx context.Context, db *sql.DB) error {
	util.MustExec(db, "CREATE TABLE IF NOT EXISTS reverse_scan (id BIGINT PRIMARY KEY NOT NULL, balance BIGINT NOT NULL, pad VARCHAR(128))")

	r := rand.New(rand.NewSource(time.Now().Unix()))
	batchSize := 100
	offset := 0
	args := make([]string, batchSize)
	if c.totalRows%batchSize != 0 {
		c.totalRows = (c.totalRows / batchSize) * batchSize
	}
	for offset < c.totalRows {
		for i := 0; i < batchSize; i++ {
			len := r.Intn(128)
			str := make([]byte, len+1)
			util.RandString(str, r)
			args[i] = fmt.Sprintf("(%d, 1000, \"%s\")", offset+i, str)
		}

		query := fmt.Sprintf("INSERT IGNORE INTO reverse_scan (id, balance, pad) values %s", strings.Join(args, ","))
		insertF := func() error {
			_, err := db.Exec(query)
			return err
		}
		err := util.RunWithRetry(ctx, 10, 3*time.Second, insertF)
		if err != nil {
			client.Fatalf("exec %s err %s", query, err)
		}

		client.Infof("[%s] insert %d rows", c, batchSize)

		offset += batchSize
	}

	c.reverseScanSum(db)
	return nil
}

func (c *ReverseScanCase) Execute(ctx context.Context, db *sql.DB) error {
	client.Infof("[%s] start to test...", c)
	defer func() {
		client.Infof("[%s] test end...", c)
	}()

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		// Transfer money
		c.transferMoneySeveralTimes(ctx, db, 1000)

		// Using reverse scan to sum the balance
		time.Sleep(5 * time.Second)
		c.reverseScanSum(db)
	}
	return nil
}

func (c *ReverseScanCase) reverseScanSum(db *sql.DB) {
	var sum int
	start := time.Now()

	query := fmt.Sprintf("select sum(a.balance) from (select balance from reverse_scan order by id desc) as a")
	err := db.QueryRow(query).Scan(&sum)
	if err != nil {
		client.Errorf("[%s] select sum err %v", c, err)
		return
	}

	expectedSum := c.totalRows * 1000
	if sum != expectedSum {
		client.Fatalf("[%s] reverse_scan balance sum must be %d, but got %d", c, expectedSum, sum)
	}

	secs := time.Since(start).Seconds()
	if secs > 300 {
		client.Fatalf("[%s] reverse scan takes %f seconds to finish", c, secs)
	}
}

func (c *ReverseScanCase) transferMoneySeveralTimes(ctx context.Context, db *sql.DB, round int) {
	var wg sync.WaitGroup
	wg.Add(c.concurrency)
	for i := 0; i < c.concurrency; i++ {
		go func() {
			defer wg.Done()

			for j := 0; j < round; j++ {
				select {
				case <-ctx.Done():
					return
				default:
				}
				c.transferMoney(db)
			}
		}()
	}
	wg.Wait()
}

func (c *ReverseScanCase) transferMoney(db *sql.DB) {
	var (
		from, to int
	)

	for {
		from, to = rand.Intn(c.totalRows), rand.Intn(c.totalRows)
		if from == to {
			continue
		}
		break
	}

	amount := rand.Intn(99)
	err := c.execTxn(db, from, to, amount)
	if err != nil {
		client.Errorf("[%s] execTxn from %d to %d amount %d err %v", c, from, to, amount, err)
	}
}

func (c *ReverseScanCase) execTxn(db *sql.DB, from, to, amount int) error {
	tx, err := db.Begin()
	if err != nil {
		return errors.Trace(err)
	}
	defer tx.Rollback()

	rows, err := tx.Query(fmt.Sprintf("SELECT id, balance FROM reverse_scan WHERE id in (%d, %d) FOR UPDATE", from, to))
	if err != nil {
		return errors.Trace(err)
	}
	defer rows.Close()

	var (
		fromBalance int
		toBalance   int
		count       int = 0
	)
	for rows.Next() {
		var id, balance int
		if err = rows.Scan(&id, &balance); err != nil {
			return errors.Trace(err)
		}
		switch id {
		case from:
			fromBalance = balance
		case to:
			toBalance = balance
		default:
			client.Fatalf("[%s] got unexpected acount %d", c, id)
		}

		count += 1
	}
	if err = rows.Err(); err != nil {
		return errors.Trace(err)
	}
	if count != 2 {
		client.Fatal("[%s] select %d(%d) -> %d(%d) invalid count %d", c, from, fromBalance, to, toBalance, count)
	}

	if fromBalance >= amount {
		update := fmt.Sprintf(`
		UPDATE reverse_scan
		SET balance = CASE id WHEN %d THEN %d WHEN %d THEN %d END
		WHERE id IN (%d, %d)`, to, toBalance+amount, from, fromBalance-amount, from, to)
		_, err = tx.Exec(update)
		if err != nil {
			return errors.Trace(err)
		}
	}

	return errors.Trace(tx.Commit())
}