// 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 core

import (
	"context"

	"github.com/pingcap/parser/ast"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/types"
)

type gcSubstituter struct {
}

// ExprColumnMap is used to store all expressions of indexed generated columns in a table,
// and map them to the generated columns,
// thus we can substitute the expression in a query to an indexed generated column.
type ExprColumnMap map[expression.Expression]*expression.Column

// optimize try to replace the expression to indexed virtual generate column in where, group by, order by, and field clause
// so that we can use the index on expression.
// For example: select a+1 from t order by a+1, with a virtual generate column c as (a+1) and
// an index on c. We need to replace a+1 with c so that we can use the index on c.
// See also https://dev.mysql.com/doc/refman/8.0/en/generated-column-index-optimizations.html
func (gc *gcSubstituter) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
	exprToColumn := make(ExprColumnMap)
	collectGenerateColumn(lp, exprToColumn)
	if len(exprToColumn) == 0 {
		return lp, nil
	}
	return gc.substitute(ctx, lp, exprToColumn), nil
}

// Only integers can be substituted between different types now.
// Because they are all stored as 64-bit integers in chunks.
// And the inserted value will meet the requirements of virtual column type and virtual column expression type.
func getSubstitutableType(colType *types.FieldType) int {
	switch colType.Tp {
	case mysql.TypeTiny, mysql.TypeShort, mysql.TypeInt24, mysql.TypeLong, mysql.TypeLonglong:
		return -1
	}
	return int(colType.Tp)
}

// We do not check whether flen is the same.
// Because different flen will not cause data truncation.
// See https://github.com/pingcap/tidb/pull/16316.
func checkSubstitutability(c1, c2 *types.FieldType) bool {
	partialEqual := c1.Decimal == c2.Decimal &&
		c1.Charset == c2.Charset &&
		c1.Collate == c2.Collate &&
		getSubstitutableType(c1) == getSubstitutableType(c2) &&
		mysql.HasUnsignedFlag(c1.Flag) == mysql.HasUnsignedFlag(c2.Flag)
	if !partialEqual || len(c1.Elems) != len(c2.Elems) {
		return false
	}
	for i := range c1.Elems {
		if c1.Elems[i] != c2.Elems[i] {
			return false
		}
	}
	return true
}

// collectGenerateColumn collect the generate column and save them to a map from their expressions to themselves.
// For the sake of simplicity, we don't collect the stored generate column because we can't get their expressions directly.
// TODO: support stored generate column.
func collectGenerateColumn(lp LogicalPlan, exprToColumn ExprColumnMap) {
	for _, child := range lp.Children() {
		collectGenerateColumn(child, exprToColumn)
	}
	ds, ok := lp.(*DataSource)
	if !ok {
		return
	}
	tblInfo := ds.tableInfo
	for _, idx := range tblInfo.Indices {
		for _, idxPart := range idx.Columns {
			colInfo := tblInfo.Columns[idxPart.Offset]
			if colInfo.IsGenerated() && !colInfo.GeneratedStored {
				s := ds.schema.Columns
				col := expression.ColInfo2Col(s, colInfo)
				if col != nil && checkSubstitutability(col.GetType(), col.VirtualExpr.GetType()) {
					exprToColumn[col.VirtualExpr] = col
				}
			}
		}
	}
}

func tryToSubstituteExpr(expr *expression.Expression, sctx sessionctx.Context, candidateExpr expression.Expression, tp types.EvalType, schema *expression.Schema, col *expression.Column) {
	if (*expr).Equal(sctx, candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
		schema.ColumnIndex(col) != -1 {
		*expr = col
	}
}

func (gc *gcSubstituter) substitute(ctx context.Context, lp LogicalPlan, exprToColumn ExprColumnMap) LogicalPlan {
	sctx := lp.SCtx().GetSessionVars().StmtCtx
	var expr *expression.Expression
	var tp types.EvalType
	switch x := lp.(type) {
	case *LogicalSelection:
		for _, cond := range x.Conditions {
			sf, ok := cond.(*expression.ScalarFunction)
			if !ok {
				continue
			}
			switch sf.FuncName.L {
			case ast.EQ, ast.LT, ast.LE, ast.GT, ast.GE:
				if sf.GetArgs()[0].ConstItem(sctx) {
					expr = &sf.GetArgs()[1]
					tp = sf.GetArgs()[0].GetType().EvalType()
				} else if sf.GetArgs()[1].ConstItem(sctx) {
					expr = &sf.GetArgs()[0]
					tp = sf.GetArgs()[1].GetType().EvalType()
				} else {
					continue
				}
				for candidateExpr, column := range exprToColumn {
					tryToSubstituteExpr(expr, lp.SCtx(), candidateExpr, tp, x.Schema(), column)
				}
			case ast.In:
				expr = &sf.GetArgs()[0]
				tp = sf.GetArgs()[1].GetType().EvalType()
				canSubstitute := true
				// Can only substitute if all the operands on the right-hand
				// side are constants of the same type.
				for i := 1; i < len(sf.GetArgs()); i++ {
					if !sf.GetArgs()[i].ConstItem(sctx) || sf.GetArgs()[i].GetType().EvalType() != tp {
						canSubstitute = false
						break
					}
				}
				if canSubstitute {
					for candidateExpr, column := range exprToColumn {
						tryToSubstituteExpr(expr, lp.SCtx(), candidateExpr, tp, x.Schema(), column)
					}
				}
			}
		}
	case *LogicalProjection:
		for i := range x.Exprs {
			tp = x.Exprs[i].GetType().EvalType()
			for candidateExpr, column := range exprToColumn {
				tryToSubstituteExpr(&x.Exprs[i], lp.SCtx(), candidateExpr, tp, x.children[0].Schema(), column)
			}
		}
	case *LogicalSort:
		for i := range x.ByItems {
			tp = x.ByItems[i].Expr.GetType().EvalType()
			for candidateExpr, column := range exprToColumn {
				tryToSubstituteExpr(&x.ByItems[i].Expr, lp.SCtx(), candidateExpr, tp, x.Schema(), column)
			}
		}
		// TODO: Uncomment these code after we support virtual generate column push down.
		//case *LogicalAggregation:
		//	for _, aggFunc := range x.AggFuncs {
		//		for i := 0; i < len(aggFunc.Args); i++ {
		//			tp = aggFunc.Args[i].GetType().EvalType()
		//			for candidateExpr, column := range exprToColumn {
		//				if aggFunc.Args[i].Equal(lp.SCtx(), candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
		//					x.Schema().ColumnIndex(column) != -1 {
		//					aggFunc.Args[i] = column
		//				}
		//			}
		//		}
		//	}
		//	for i := 0; i < len(x.GroupByItems); i++ {
		//		tp = x.GroupByItems[i].GetType().EvalType()
		//		for candidateExpr, column := range exprToColumn {
		//			if x.GroupByItems[i].Equal(lp.SCtx(), candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
		//				x.Schema().ColumnIndex(column) != -1 {
		//				x.GroupByItems[i] = column
		//				x.groupByCols = append(x.groupByCols, column)
		//			}
		//		}
		//	}
	}
	for _, child := range lp.Children() {
		gc.substitute(ctx, child, exprToColumn)
	}
	return lp
}

func (*gcSubstituter) name() string {
	return "generate_column_substitute"
}