package main

import (
	"database/sql"
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/bench/bench"
	"github.com/pingcap/schrodinger-test/pkg/statdb"
	"github.com/pingcap/schrodinger-test/util"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/pkg/types"
)

const (
	blockWriteTableName   = "bench_block_write"
	blockWriteTableSchema = `(
        block_id BIGINT NOT NULL,
        block_data BLOB NOT NULL,
        PRIMARY KEY (block_id)
	)`
)

type BlockWriteConfig struct {
	Duration     time.Duration
	WarmUp       time.Duration
	NumThreads   int
	MinBlockSize int
	MaxBlockSize int
}

type BlockWrite struct {
	cfg     *BlockWriteConfig
	db      *sql.DB
	blockID uint64
}

func NewBlockWrite(cfg *BlockWriteConfig, db *sql.DB) *BlockWrite {
	return &BlockWrite{
		cfg:     cfg,
		db:      db,
		blockID: 0,
	}
}

func (b *BlockWrite) prepare() error {
	return errors.Trace(util.RecreateTable(b.db, blockWriteTableName, blockWriteTableSchema))
}

func (b *BlockWrite) run() (*types.TestResult, error) {
	rs, err := bench.ParallelSQLBench(b.name(), b.execute, b.cfg.Duration, b.cfg.NumThreads, b.db)
	if err != nil {
		client.Errorf("[bench_block_write] parallel exec failed: %v", err)
		return nil, errors.Trace(err)
	}

	cr := &types.TestResult{
		TestName: b.name(),
		QPS:      rs.Qps,
		EPS:      rs.Eps,
		AvgMs:    rs.Histogram.AvgMs,
		Pct90:    rs.Histogram.Pct90Ms,
		Pct99:    rs.Histogram.Pct99Ms,
	}

	return cr, nil
}

func (b *BlockWrite) execute(db *statdb.StatDB, rander *rand.Rand) error {
	blockID := atomic.AddUint64(&b.blockID, 1)
	blockSize := rander.Intn(b.cfg.MaxBlockSize-b.cfg.MinBlockSize) + b.cfg.MinBlockSize

	blockData := util.RandomASCIIBytes(rander, blockSize)

	stmt := fmt.Sprintf(
		"INSERT INTO `%s` (`block_id`, `block_data`) VALUES (?, ?)", blockWriteTableName)
	_, err := db.Exec(stmt, blockID, blockData)
	return errors.Trace(err)
}

func (b *BlockWrite) name() string {
	return "bench_block_write"
}