// Copyright 2019 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 randsql

import (
	"fmt"
	"math/rand"

	. "github.com/pingcap/schrodinger-test/logic/ddl"
)

func (c *RandSQLGen) GenDML(tp DMLKind) *DMLJobTask {
	switch tp {
	case DMLInsert:
		return c.GenRandInsert()
	case DMLUpdate:
		return c.GenRandUpdate()
	case DMLDelete:
		return c.GenRandDelete()
	case DMLReplace:
		return c.GenRandReplace()
	}
	return nil
}

func (c *RandSQLGen) DoDMLJobInLocal(task *DMLJobTask) error {
	switch task.k {
	case DMLInsert:
		return c.doInsertJob(task)
	case DMLUpdate:
		return c.doUpdateJob(task)
	case DMLDelete:
		return c.doDeleteJob(task)
	case DMLReplace:
		return c.doReplaceJob(task)
	}
	return fmt.Errorf("unknow dml task , %v", *task)
}

const (
	DMLInsert DMLKind = iota
	DMLUpdate
	DMLDelete
	DMLReplace
	DMLSelect
)

type SQLTask interface {
	SQL() string
	Table() *TableInfo
}

type DMLJobTask struct {
	k            DMLKind
	tblInfo      *TableInfo
	sql          string
	assigns      []*ColumnDescriptor
	whereColumns []*ColumnDescriptor
	onDupRowIdx  int
	err          error
}

func (task *DMLJobTask) SQL() string {
	return task.sql
}

func (task *DMLJobTask) Table() *TableInfo {
	return task.tblInfo
}

func (task *DMLJobTask) GetTableName() string {
	return task.tblInfo.name
}

type InsertColumnStrategy = int
type InsertMissingValueStrategy = int

const (
	InsertColumnStrategyZeroNonPk InsertColumnStrategy = iota
	InsertColumnStrategyAllNonPk
	InsertColumnStrategyRandomNonPk
	InsertColumnStrategyEnd
)

const (
	InsertMissingValueStrategyAllNull InsertMissingValueStrategy = iota
	InsertMissingValueStrategyAllDefault
	InsertMissingValueStrategyRandom
	InsertMissingValueStrategyEnd
)

type ddlTestInsertConfig struct {
	columnStrategy       InsertColumnStrategy       // how non-Primary-Key columns are picked
	missingValueStrategy InsertMissingValueStrategy // how columns are filled when they are not picked in VALUE statement
}

func (c *RandSQLGen) GenRandInsert() *DMLJobTask {
	config := ddlTestInsertConfig{
		columnStrategy:       rand.Intn(InsertColumnStrategyEnd),
		missingValueStrategy: rand.Intn(InsertMissingValueStrategyEnd),
	}
	return c.GenInsert(config)
}

func (c *RandSQLGen) prepareInsert(config ddlTestInsertConfig, onDup bool) (*TableInfo, []*ColumnInfo, []*ColumnDescriptor, int) {
	c.tablesLock.RLock()
	table := c.pickupRandomTable()
	if table == nil {
		c.tablesLock.RUnlock()
		return nil, nil, nil, -1
	}
	table.lock.RLock()
	columns := table.filterColumns(table.predicateNotGenerated)
	nonPkColumns := table.filterColumns(table.predicateNonPrimaryKeyAndNotGen)
	table.lock.RUnlock()
	c.tablesLock.RUnlock()

	onDupRowIdx := -1
	if onDup && (rand.Float64() >= 0.5) && len(columns[0].rows) > 0 {
		onDupRowIdx = rand.Intn(len(columns[0].rows))
	}

	// build assignments
	assigns := make([]*ColumnDescriptor, 0)
	for _, column := range columns {
		pick := false
		if column.isPrimaryKey {
			// PrimaryKey Column is always assigned values
			pick = true
		} else {
			// NonPrimaryKey Column is assigned by strategy
			switch config.columnStrategy {
			case InsertColumnStrategyAllNonPk:
				pick = true
			case InsertColumnStrategyZeroNonPk:
				pick = false
			case InsertColumnStrategyRandomNonPk:
				if rand.Float64() <= float64(1)/float64(len(nonPkColumns)) {
					pick = true
				}
			}
		}
		if pick {
			// check unique value when inserting into a column of primary key
			if column.isPrimaryKey {
				if onDupRowIdx >= 0 {
					assigns = append(assigns, &ColumnDescriptor{column, column.rows[onDupRowIdx]})
				} else {
					if newValue, ok := column.randValueUnique(column.rows); ok {
						assigns = append(assigns, &ColumnDescriptor{column, newValue})
					} else {
						return nil, nil, nil, -1
					}
				}
			} else {
				assigns = append(assigns, &ColumnDescriptor{column, column.randValue()})
			}
		}
	}
	return table, columns, assigns, onDupRowIdx

}

func (c *RandSQLGen) GenInsert(cfg interface{}) *DMLJobTask {
	config := cfg.(ddlTestInsertConfig)
	table, _, assigns, onDupRowIdx := c.prepareInsert(config, true)
	sql := ""
	if len(assigns) == 0 {
		return nil
	}
	sql = fmt.Sprintf("INSERT INTO `%s` SET ", table.name)
	perm := rand.Perm(len(assigns))
	for i, idx := range perm {
		assign := assigns[idx]
		if i > 0 {
			sql += ", "
		}
		sql += fmt.Sprintf("`%s` = %v", assign.column.name, assign.getValueString())
	}

	if onDupRowIdx >= 0 {
		updatePk := rand.Float64() >= 0.5
		sql += "ON DUPLICATE KEY UPDATE "
		for i, assign := range assigns {
			if i > 0 {
				sql += ", "
			}
			if assign.column.isPrimaryKey && updatePk {
				if newValue, ok := assign.column.randValueUnique(assign.column.rows); ok {
					assigns[i].value = newValue
				} else {
					return nil
				}
			}
			sql += fmt.Sprintf("`%s` = %v", assign.column.name, assign.getValueString())
		}
	}

	task := &DMLJobTask{
		k:           DMLInsert,
		sql:         sql,
		tblInfo:     table,
		assigns:     assigns,
		onDupRowIdx: onDupRowIdx,
	}
	return task
}

func (c *RandSQLGen) doInsertOnlyJob(task *DMLJobTask) error {
	table := task.tblInfo
	assigns := task.assigns

	// append row
	table.lock.Lock()
	for _, column := range table.columns {
		cd := column.getMatchedColumnDescriptor(assigns)
		if cd == nil {
			if column.isGenerated() {
				cd = column.dependency.getMatchedColumnDescriptor(assigns)
				if cd == nil {
					column.rows = append(column.rows, nil)
				} else {
					column.rows = append(column.rows, cd.column.getDependenciedColsValue(column))
				}
			} else {
				// only happens when using SET
				column.rows = append(column.rows, column.defaultValue)
			}
		} else {
			column.rows = append(column.rows, cd.value)
		}
	}
	table.numberOfRows++
	table.lock.Unlock()
	return nil
}

func (c *RandSQLGen) doInsertJob(task *DMLJobTask) error {
	if task.onDupRowIdx < 0 {
		return c.doInsertOnlyJob(task)
	}
	table := task.tblInfo
	assigns := task.assigns

	// update row
	table.lock.Lock()
	i := task.onDupRowIdx
	for _, cd := range assigns {
		cd.column.rows[i] = cd.value
		if cd.column.hasGenerateCol() {
			for _, col := range cd.column.dependenciedCols {
				col.rows[i] = cd.column.getDependenciedColsValue(col)
			}
		}
	}
	table.lock.Unlock()
	return nil
}

func (c *RandSQLGen) GenRandReplace() *DMLJobTask {
	config := ddlTestInsertConfig{
		columnStrategy:       rand.Intn(InsertColumnStrategyEnd),
		missingValueStrategy: rand.Intn(InsertMissingValueStrategyEnd),
	}
	return c.GenReplace(config)
}

func (c *RandSQLGen) GenReplace(config ddlTestInsertConfig) *DMLJobTask {
	table, _, assigns, onDupRowIdx := c.prepareInsert(config, true)
	// build SQL
	sql := ""
	if len(assigns) == 0 {
		return nil
	}
	sql = fmt.Sprintf("REPLACE INTO `%s` SET ", table.name)
	perm := rand.Perm(len(assigns))
	for i, idx := range perm {
		assign := assigns[idx]
		if i > 0 {
			sql += ", "
		}
		sql += fmt.Sprintf("`%s` = %v", assign.column.name, assign.getValueString())
	}

	task := &DMLJobTask{
		k:           DMLReplace,
		sql:         sql,
		tblInfo:     table,
		assigns:     assigns,
		onDupRowIdx: onDupRowIdx,
	}
	return task
}

func (c *RandSQLGen) doReplaceJob(task *DMLJobTask) error {
	return c.doInsertJob(task)
}

type WhereStrategy = int

const (
	WhereStrategyNone WhereStrategy = iota
	WhereStrategyRandomInPk
	WhereStrategyRandomInNonPk
	WhereStrategyRandomMixed
	WhereStrategyEnd
)

func randWhereStrategy() WhereStrategy {
	p := rand.Float64()
	if p < 0.3 {
		return WhereStrategyRandomInPk
	} else if p < 0.6 {
		return WhereStrategyRandomInNonPk
	} else if p < 0.9 {
		return WhereStrategyRandomMixed
	}
	return WhereStrategyNone
}

type UpdateTargetStrategy = int

const (
	UpdateTargetStrategyAllColumns UpdateTargetStrategy = iota
	UpdateTargetStrategyRandom
	UpdateTargetStrategyEnd
)

type UpdateConfig struct {
	whereStrategy  WhereStrategy        // how "where" statement is generated
	targetStrategy UpdateTargetStrategy // which column to update
}

func (c *RandSQLGen) GenRandUpdate() *DMLJobTask {
	config := UpdateConfig{
		whereStrategy:  randWhereStrategy(),
		targetStrategy: rand.Intn(UpdateTargetStrategyEnd),
	}
	return c.GenUpdate(config)

}

func (c *RandSQLGen) buildWhereColumns(whereStrategy WhereStrategy, pkColumns, nonPkColumns []*ColumnInfo, numberOfRows int) []*ColumnDescriptor {
	// build where conditions
	whereColumns := make([]*ColumnDescriptor, 0)
	if whereStrategy == WhereStrategyRandomInPk || whereStrategy == WhereStrategyRandomMixed {
		if len(pkColumns) > 0 {
			picks := rand.Intn(len(pkColumns))
			perm := rand.Perm(picks)
			for _, idx := range perm {
				// value will be filled later
				whereColumns = append(whereColumns, &ColumnDescriptor{pkColumns[idx], -1})
			}
		}
	}
	if whereStrategy == WhereStrategyRandomInNonPk || whereStrategy == WhereStrategyRandomMixed {
		if len(nonPkColumns) > 0 {
			picks := rand.Intn(len(nonPkColumns))
			perm := rand.Perm(picks)
			for _, idx := range perm {
				// value will be filled later
				whereColumns = append(whereColumns, &ColumnDescriptor{nonPkColumns[idx], -1})
			}
		}
	}

	// fill values of where statements
	if len(whereColumns) > 0 {
		rowToUpdate := rand.Intn(numberOfRows)
		for _, cd := range whereColumns {
			cd.value = cd.column.rows[rowToUpdate]
		}
	}

	return whereColumns
}

func (c *RandSQLGen) GenUpdate(config UpdateConfig) *DMLJobTask {
	c.tablesLock.RLock()
	table := c.pickupRandomTable()
	if table == nil {
		c.tablesLock.RUnlock()
		return nil
	}
	table.lock.RLock()
	pkColumns := table.filterColumns(table.predicatePrimaryKey)
	nonPkColumnsAndCanBeWhere := table.filterColumns(table.predicateNonPrimaryKeyAndCanBeWhere)
	nonPkColumnsAndNotGen := table.filterColumns(table.predicateNonPrimaryKeyAndNotGen)
	table.lock.RUnlock()
	c.tablesLock.RUnlock()

	if table.numberOfRows == 0 {
		return nil
	}

	// build where conditions
	whereColumns := c.buildWhereColumns(config.whereStrategy, pkColumns, nonPkColumnsAndCanBeWhere, table.numberOfRows)

	// build assignments
	assigns := make([]*ColumnDescriptor, 0)
	picks := 0
	switch config.targetStrategy {
	case UpdateTargetStrategyRandom:
		if len(nonPkColumnsAndNotGen) > 0 {
			picks = rand.Intn(len(nonPkColumnsAndNotGen))
		}
	case UpdateTargetStrategyAllColumns:
		picks = len(nonPkColumnsAndNotGen)
	}
	if picks == 0 {
		return nil
	}
	perm := rand.Perm(picks)
	for _, idx := range perm {
		assigns = append(assigns, &ColumnDescriptor{nonPkColumnsAndNotGen[idx], nonPkColumnsAndNotGen[idx].randValue()})
	}

	// build SQL
	sql := fmt.Sprintf("UPDATE `%s` SET ", table.name)
	for i, cd := range assigns {
		if i > 0 {
			sql += ", "
		}
		sql += fmt.Sprintf("`%s` = %v", cd.column.name, cd.getValueString())
	}
	if len(whereColumns) > 0 {
		sql += " WHERE "
		for i, cd := range whereColumns {
			if i > 0 {
				sql += " AND "
			}
			sql += cd.buildConditionSQL()
		}
	}

	task := &DMLJobTask{
		k:            DMLUpdate,
		tblInfo:      table,
		sql:          sql,
		assigns:      assigns,
		whereColumns: whereColumns,
	}
	return task
}

func (c *RandSQLGen) doUpdateJob(task *DMLJobTask) error {
	table := task.tblInfo
	assigns := task.assigns
	whereColumns := task.whereColumns

	// update values
	table.lock.RLock()
	for i := 0; i < table.numberOfRows; i++ {
		match := true
		for _, cd := range whereColumns {
			if cd.value != cd.column.rows[i] {
				match = false
				break
			}
		}
		if match {
			for _, cd := range assigns {
				cd.column.rows[i] = cd.value
				if cd.column.hasGenerateCol() {
					for _, col := range cd.column.dependenciedCols {
						col.rows[i] = cd.column.getDependenciedColsValue(col)
					}
				}
			}
		}
	}
	table.lock.RUnlock()
	return nil

}

type DeleteConfig struct {
	whereStrategy WhereStrategy // how "where" statement is generated
}

func (c *RandSQLGen) GenRandDelete() *DMLJobTask {
	config := DeleteConfig{
		whereStrategy: randWhereStrategy(),
	}
	return c.GenDelete(config)
}

func (c *RandSQLGen) GenDelete(config DeleteConfig) *DMLJobTask {
	c.tablesLock.RLock()
	table := c.pickupRandomTable()
	if table == nil {
		c.tablesLock.RUnlock()
		return nil
	}
	table.lock.RLock()
	pkColumns := table.filterColumns(table.predicatePrimaryKey)
	nonPkColumnsAndCanBeWhere := table.filterColumns(table.predicateNonPrimaryKeyAndCanBeWhere)
	table.lock.RUnlock()
	c.tablesLock.RUnlock()

	if table.numberOfRows == 0 {
		return nil
	}

	whereColumns := c.buildWhereColumns(config.whereStrategy, pkColumns, nonPkColumnsAndCanBeWhere, table.numberOfRows)

	// build SQL
	sql := fmt.Sprintf("DELETE FROM `%s`", table.name)
	if len(whereColumns) > 0 {
		sql += " WHERE "
		for i, cd := range whereColumns {
			if i > 0 {
				sql += " AND "
			}
			sql += cd.buildConditionSQL()
		}
	}

	task := &DMLJobTask{
		k:            DMLDelete,
		tblInfo:      table,
		sql:          sql,
		whereColumns: whereColumns,
	}
	return task
}

func (c *RandSQLGen) doDeleteJob(task *DMLJobTask) error {
	table := task.tblInfo
	whereColumns := task.whereColumns

	// update values
	table.lock.Lock()
	for i := table.numberOfRows - 1; i >= 0; i-- {
		match := true
		for _, cd := range whereColumns {
			if cd.value != cd.column.rows[i] {
				match = false
				break
			}
		}
		if match {
			// we must use `table.columns` here, since there might be new columns after deletion
			for _, column := range table.columns {
				column.rows = append(column.rows[:i], column.rows[i+1:]...)
			}
			table.numberOfRows--
		}
	}
	table.lock.Unlock()

	return nil
}

func (c *RandSQLGen) GenRandSelect() *DMLJobTask {
	return c.GenSelect(randWhereStrategy())
}

func (c *RandSQLGen) GenSelect(ws WhereStrategy) *DMLJobTask {
	c.tablesLock.RLock()
	table := c.pickupRandomTable()
	if table == nil {
		c.tablesLock.RUnlock()
		return nil
	}
	c.tablesLock.RUnlock()
	return c.genSelectWithTable(ws, table)
}

func (c *RandSQLGen) GenRandSelectWithTableName(name string) *DMLJobTask {
	return c.GenSelectWithTableName(randWhereStrategy(), name)
}

func (c *RandSQLGen) GenSelectWithTableName(ws WhereStrategy, tableName string) *DMLJobTask {
	c.tablesLock.RLock()
	table := c.tables[tableName]
	if table == nil {
		c.tablesLock.RUnlock()
		return nil
	}
	c.tablesLock.RUnlock()
	return c.genSelectWithTable(ws, table)
}

func getRandSelectPartitionNames(table *TableInfo) string {
	if table.partition == nil {
		return ""
	}
	return fmt.Sprintf("partition ( %v )", table.pickRandPartitionNames())
}

func (c *RandSQLGen) genSelectWithTable(ws WhereStrategy, table *TableInfo) *DMLJobTask {
	table.lock.RLock()
	pkColumns := table.filterColumns(table.predicatePrimaryKey)
	nonPkColumnsAndCanBeWhere := table.filterColumns(table.predicateNonPrimaryKeyAndCanBeWhere)
	table.lock.RUnlock()

	if table.numberOfRows == 0 {
		return nil
	}

	// build SQL
	sql := fmt.Sprintf("SELECT * FROM `%s`", table.name)

	// select table with partition.
	if table.isPartitionTable() && rand.Float64() > 0.5 {
		sql += getRandSelectPartitionNames(table)
	}

	// build where conditions
	whereColumns := c.buildWhereColumns(ws, pkColumns, nonPkColumnsAndCanBeWhere, table.numberOfRows)

	if len(whereColumns) > 0 {
		sql += " WHERE "
		for i, cd := range whereColumns {
			if i > 0 {
				sql += " AND "
			}
			sql += cd.buildConditionSQL()
		}
	}

	task := &DMLJobTask{
		k:            DMLSelect,
		tblInfo:      table,
		sql:          sql,
		whereColumns: whereColumns,
	}
	return task
}

func (c *RandSQLGen) GenAdminCheck() *DMLJobTask {
	if len(c.tables) == 0 {
		return nil
	}

	sql := "ADMIN CHECK TABLE "
	i := 0
	for _, table := range c.tables {
		if i > 0 {
			sql += ", "
		}
		sql += fmt.Sprintf("`%s`", table.name)
		i++
	}
	task := &DMLJobTask{
		sql: sql,
	}
	return task
}