package oas

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/opcode"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

var (
	// ExplainFormatBrief is the same as ExplainFormatRow, but it ignores explain ID suffix.
	ExplainFormatBrief = "brief"
	// ExplainFormatDOT indicates that using DotGraph to explain.
	ExplainFormatDOT = "dot"
	// ExplainFormatHint is to show hint information.
	ExplainFormatHint = "hint"
	// ExplainFormatJSON displays the information in JSON format.
	ExplainFormatJSON = "json"
	// ExplainFormatROW presents the output in tabular format.
	ExplainFormatROW = "row"
	// ExplainFormatVerbose display additional information regarding the plan.
	ExplainFormatVerbose = "verbose"
	// ExplainFormatTraditional is the same as ExplainFormatROW.
	ExplainFormatTraditional = "traditional"
	// ExplainFormatTrueCardCost indicates the optimizer to use true cardinality to calculate the cost.
	ExplainFormatTrueCardCost = "true_card_cost"
	// ExplainFormatBinary prints the proto for binary plan.
	ExplainFormatBinary = "binary"
	// ExplainFormatTiDBJSON warp the default result in JSON format
	ExplainFormatTiDBJSON = "tidb_json"
	// ExplainFormatCostTrace prints the cost and cost formula of each operator.
	ExplainFormatCostTrace = "cost_trace"
	// ExplainFormatPlanCache prints the reason why can't use non-prepared plan cache by warning
	ExplainFormatPlanCache = "plan_cache"

	// ExplainFormats stores the valid formats for explain statement, used by validator.
	ExplainFormats = []string{
		ExplainFormatBrief,
		ExplainFormatDOT,
		ExplainFormatHint,
		ExplainFormatJSON,
		ExplainFormatROW,
		ExplainFormatVerbose,
		ExplainFormatTraditional,
		ExplainFormatTrueCardCost,
		ExplainFormatBinary,
		ExplainFormatTiDBJSON,
		ExplainFormatCostTrace,
		ExplainFormatPlanCache,
	}
)

// query result of infomation_schema.columns
type MetaColumn struct {
	TableSchema string `json:"table_schema"`
	TableName   string `json:"table_name"`
	ColumnName  string `json:"column_name"`
	DataType    string `json:"data_type"`
	IsNullable  string `json:"is_nullable"`
}

type Column struct {
	ColumnName string `json:"column_name"`
	Database   string `json:"-"`
	Table      string `json:"-"`
	// currently we can decide the SQL return type since some functions may changge the origin column type
	// DataType   string `json:"data_type"`
	// Nullable   bool   `json:"is_nullable"`
}

type Columns []*Column

type Tables map[string]Columns

type Databases map[string]Tables

type CTE struct {
	Def     *ast.CommonTableExpression `json:"def"`
	Columns Columns                    `json:"columns"`
}

type Extractor struct {
	databases Databases
	currentDB string
	ctes      []CTE
}

func NewExtractor(metaColumns []*MetaColumn) *Extractor {
	databases := make(Databases)
	for _, mc := range metaColumns {
		tables, ok := databases[mc.TableSchema]
		if !ok {
			tables = make(Tables)
		}
		columns, ok := tables[mc.TableName]
		if !ok {
			columns = make(Columns, 0, 4)
		}
		columns = append(columns, &Column{
			ColumnName: mc.ColumnName,
			Table:      mc.TableName,
			Database:   mc.TableSchema,
		})
		tables[mc.TableName] = columns
		databases[mc.TableSchema] = tables
	}
	return &Extractor{databases: databases}
}

func (e *Extractor) Extract(ctx context.Context, sql string) (Columns, error) {
	p := parser.New()
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return nil, err
	}

	result := make(Columns, 0)
	last := len(stmtNodes) - 1
	for i, node := range stmtNodes {
		if i != last && reflect.TypeOf(node).String() != "*ast.UseStmt" {
			tilog.Info(ctx, "skip statement since not last statement", zap.Any("type", reflect.TypeOf(node)), zap.String("stmt", node.Text()))
			continue
		}
		switch x := node.(type) {
		case *ast.SelectStmt:
			return e.extractSelect(x)
		case *ast.UseStmt:
			e.setCurrentDB(x.DBName)
		case *ast.SetOprStmt:
			return e.extractSetOpr(x)
		case *ast.ShowStmt:
			return e.extractShow(x)
		case *ast.ExplainStmt:
			return e.extractExplain(x)
		default:
			tilog.Info(ctx, "skip statement since unknown node type", zap.Any("node", x), zap.String("type", fmt.Sprintf("%T", x)), zap.String("stmt", node.Text()))
		}
	}
	return result, nil
}

func (e *Extractor) setCurrentDB(dbName string) {
	e.currentDB = strings.ToLower(dbName)
}

func (e *Extractor) extractSelect(sel *ast.SelectStmt) (Columns, error) {
	if sel.With != nil {
		l := len(e.ctes)
		defer func() {
			e.ctes = e.ctes[:l]
		}()
		err := e.extractWith(sel.With)
		if err != nil {
			return nil, err
		}
	}

	columns, err := e.extractTableRefs(sel.From)
	if err != nil {
		return nil, err
	}
	// tilog.Info(context.Background(), "columns", zap.Any("fields", sel.Fields.Fields))
	columns, err = e.extractProjection(sel.Fields.Fields, columns)
	if err != nil {
		return nil, err
	}
	return columns, nil
}

func getInnerFromParenthesesAndUnaryPlus(expr ast.ExprNode) ast.ExprNode {
	if pexpr, ok := expr.(*ast.ParenthesesExpr); ok {
		return getInnerFromParenthesesAndUnaryPlus(pexpr.Expr)
	}
	if uexpr, ok := expr.(*ast.UnaryOperationExpr); ok && uexpr.Op == opcode.Plus {
		return getInnerFromParenthesesAndUnaryPlus(uexpr.V)
	}
	return expr
}

func (e *Extractor) extractProjection(fields []*ast.SelectField, oldCols Columns) (Columns, error) {
	newCols := make(Columns, 0, len(fields))
	for _, field := range fields {
		if field.WildCard != nil {
			dbName := field.WildCard.Schema
			tblName := field.WildCard.Table
			for _, oldCol := range oldCols {
				if (dbName.L == "" || dbName.L == strings.ToLower(oldCol.Database)) &&
					(tblName.L == "" || tblName.L == strings.ToLower(oldCol.Table)) {
					newCols = append(newCols, oldCol)
				}
			}

		} else {
			var colName string
			innerNode := getInnerFromParenthesesAndUnaryPlus(field.Expr)
			if colNameField, ok := innerNode.(*ast.ColumnNameExpr); ok {
				if field.AsName.L != "" {
					colName = field.AsName.L
				} else {
					colName = colNameField.Name.Name.L
				}
				newCols = append(newCols, &Column{ColumnName: colName})
			} else if field.AsName.L != "" {
				newCols = append(newCols, &Column{ColumnName: field.AsName.L})
			} else {
				fieldName := field.Text()
				fieldName = strings.TrimLeft(fieldName, "\t\n ")
				fieldName = strings.TrimRight(fieldName, "\t\n ")
				fieldName = strings.ReplaceAll(fieldName, "\"", "")
				newCols = append(newCols, &Column{ColumnName: fieldName})
			}
		}
	}
	return newCols, nil
}

func (e *Extractor) extractTableRefs(from *ast.TableRefsClause) (Columns, error) {
	if from == nil {
		result := make(Columns, 0)
		return result, nil
	}
	return e.extractResultSet(from.TableRefs)
}

func (e *Extractor) extractResultSet(node ast.ResultSetNode) (Columns, error) {
	switch x := node.(type) {
	case *ast.Join:
		return e.extractJoin(x)
	case *ast.TableSource:
		var columns Columns
		var err error
		switch v := x.Source.(type) {
		case *ast.SelectStmt:
			columns, err = e.extractSelect(v)
		case *ast.SetOprStmt:
			columns, err = e.extractSetOpr(v)
		case *ast.TableName:
			columns, err = e.extractDataSource(v)
		default:
			err = fmt.Errorf("extract result set encounter unknown table source type: %T", v)
		}
		if err != nil {
			return nil, err
		}
		if x.AsName.L != "" {
			for _, column := range columns {
				column.Table = x.AsName.L
			}
		}
		return columns, nil
	case *ast.SelectStmt:
		return e.extractSelect(x)
	case *ast.SetOprStmt:
		return e.extractSetOpr(x)
	default:
		return nil, fmt.Errorf("extract result set encounter unknown result set type: %T", x)
	}
}

func (e *Extractor) extractJoin(join *ast.Join) (Columns, error) {
	if join.Right == nil {
		return e.extractResultSet(join.Left)
	}
	leftRes, err := e.extractResultSet(join.Left)
	if err != nil {
		return nil, err
	}
	rightRes, err := e.extractResultSet(join.Right)
	if err != nil {
		return nil, err
	}
	result := append(leftRes, rightRes...)
	return result, nil
}

func (e *Extractor) extractDataSource(tableName *ast.TableName) (Columns, error) {
	if tableName.Schema.L == "" {
		cols, err := e.tryExtractCTE(tableName)
		if err != nil || cols != nil {
			return cols, err
		}
	}

	var db string
	if tableName.Schema.L != "" {
		db = tableName.Schema.L
	} else {
		db = e.currentDB
	}
	if db == "" {
		return nil, errors.New("no database selected")
	}

	database, ok := e.databases[db]
	if !ok {
		return nil, fmt.Errorf("table '%s.%s' doesn't exist", db, tableName.Name.L)
	}
	columns, ok := database[tableName.Name.L]
	if !ok {
		return nil, fmt.Errorf("table '%s.%s' doesn't exist", db, tableName.Name.L)
	}
	return columns, nil
}

func (e *Extractor) extractSetOpr(setOpr *ast.SetOprStmt) (Columns, error) {
	selects := setOpr.SelectList.Selects
	switch x := selects[0].(type) {
	case *ast.SelectStmt:
		return e.extractSelect(x)
	case *ast.SetOprStmt:
		return e.extractSetOpr(x)
	default:
		return nil, fmt.Errorf("extract set opr encounter unknown select type: %T", x)
	}
}

func (e *Extractor) extractWith(with *ast.WithClause) error {
	for _, cte := range with.CTEs {
		cols, err := e.extractCTE(cte)
		if err != nil {
			return err
		}
		e.ctes = append(e.ctes, CTE{
			Def:     cte,
			Columns: cols,
		})
	}
	return nil
}

func (e *Extractor) extractCTE(cte *ast.CommonTableExpression) (Columns, error) {
	return e.extractResultSet(cte.Query.Query)
}

func (e *Extractor) tryExtractCTE(tableName *ast.TableName) (Columns, error) {
	for i := len(e.ctes) - 1; i >= 0; i-- {
		cte := e.ctes[i]
		if cte.Def.Name.L == tableName.Name.L {
			lenCols := len(cte.Def.ColNameList)
			if lenCols > 0 {
				result := make(Columns, 0, lenCols)
				for _, c := range cte.Def.ColNameList {
					result = append(result, &Column{ColumnName: c.L})
				}
				return result, nil
			}
			return cte.Columns, nil
		}
	}
	return nil, nil
}

func (e *Extractor) extractExplain(explain *ast.ExplainStmt) (Columns, error) {
	if show, ok := explain.Stmt.(*ast.ShowStmt); ok {
		return e.extractShow(show)
	}
	return e.extractExplainPlan(explain.Format, explain.Analyze)
}

func (e *Extractor) extractExplainPlan(format string, analyze bool) (Columns, error) {
	var fieldNames []string
	format = strings.ToLower(format)
	if format == ExplainFormatTraditional {
		format = ExplainFormatROW
	}
	switch {
	case (format == ExplainFormatROW || format == ExplainFormatBrief || format == ExplainFormatPlanCache) && !analyze:
		fieldNames = []string{"id", "estRows", "task", "access object", "operator info"}
	case format == ExplainFormatVerbose:
		if analyze {
			fieldNames = []string{"id", "estRows", "estCost", "actRows", "task", "access object", "execution info", "operator info", "memory", "disk"}
		} else {
			fieldNames = []string{"id", "estRows", "estCost", "task", "access object", "operator info"}
		}
	case format == ExplainFormatTrueCardCost:
		fieldNames = []string{"id", "estRows", "estCost", "costFormula", "actRows", "task", "access object", "execution info", "operator info", "memory", "disk"}
	case format == ExplainFormatCostTrace:
		if analyze {
			fieldNames = []string{"id", "estRows", "estCost", "costFormula", "actRows", "task", "access object", "execution info", "operator info", "memory", "disk"}
		} else {
			fieldNames = []string{"id", "estRows", "estCost", "costFormula", "task", "access object", "operator info"}
		}
	case (format == ExplainFormatROW || format == ExplainFormatBrief || format == ExplainFormatPlanCache) && analyze:
		fieldNames = []string{"id", "estRows", "actRows", "task", "access object", "execution info", "operator info", "memory", "disk"}
	case format == ExplainFormatDOT:
		fieldNames = []string{"dot contents"}
	case format == ExplainFormatHint:
		fieldNames = []string{"hint"}
	case format == ExplainFormatBinary:
		fieldNames = []string{"binary plan"}
	case format == ExplainFormatTiDBJSON:
		fieldNames = []string{"TiDB_JSON"}
	default:
		return nil, fmt.Errorf("explain format '%s' is not supported now", format)
	}
	result := make(Columns, 0, 4)
	for _, name := range fieldNames {
		result = append(result, &Column{ColumnName: name})
	}
	return result, nil
}

func (e *Extractor) extractShow(show *ast.ShowStmt) (Columns, error) {
	var names []string
	switch show.Tp {
	case ast.ShowProcedureStatus, ast.ShowFunctionStatus:
		names = []string{"Db", "Name", "Type", "Definer", "Modified", "Created", "Security_type", "Comment", "character_set_client", "collation_connection", "Database Collation"}
	case ast.ShowTriggers:
		names = []string{"Trigger", "Event", "Table", "Statement", "Timing", "Created", "sql_mode", "Definer", "character_set_client", "collation_connection", "Database Collation"}
	case ast.ShowEvents:
		names = []string{"Db", "Name", "Time zone", "Definer", "Type", "Execute At", "Interval Value", "Interval Field", "Starts", "Ends", "Status", "Originator", "character_set_client", "collation_connection", "Database Collation"}
	case ast.ShowWarnings, ast.ShowErrors:
		if show.CountWarningsOrErrors {
			names = []string{"Count"}
		} else {
			names = []string{"Level", "Code", "Message"}
		}
	case ast.ShowRegions:
		names = []string{"REGION_ID", "START_KEY", "END_KEY", "LEADER_ID", "LEADER_STORE_ID", "PEERS", "SCATTERING", "WRITTEN_BYTES", "READ_BYTES", "APPROXIMATE_SIZE(MB)", "APPROXIMATE_KEYS", "SCHEDULING_CONSTRAINTS", "SCHEDULING_STATE"}
	case ast.ShowEngines:
		names = []string{"Engine", "Support", "Comment", "Transactions", "XA", "Savepoints"}
	case ast.ShowConfig:
		names = []string{"Type", "Instance", "Name", "Value"}
	case ast.ShowDatabases:
		names = []string{"Database"}
	case ast.ShowOpenTables:
		names = []string{"Database", "Table", "In_use", "Name_locked"}
	case ast.ShowTables:
		if show.DBName == "" {
			show.DBName = e.currentDB
		}
		names = []string{fmt.Sprintf("Tables_in_%s", show.DBName)}
		if show.Full {
			names = append(names, "Table_type")
		}
	case ast.ShowTableStatus:
		names = []string{"Name", "Engine", "Version", "Row_format", "Rows", "Avg_row_length",
			"Data_length", "Max_data_length", "Index_length", "Data_free", "Auto_increment",
			"Create_time", "Update_time", "Check_time", "Collation", "Checksum",
			"Create_options", "Comment"}
	case ast.ShowColumns:
		if show.Full {
			names = []string{"Field", "Type", "Collation", "Null", "Key", "Default", "Extra", "Privileges", "Comment"}
		} else {
			names = []string{"Field", "Type", "Null", "Key", "Default", "Extra"}
		}
	case ast.ShowCharset:
		names = []string{"Charset", "Description", "Default collation", "Maxlen"}
	case ast.ShowVariables, ast.ShowStatus:
		names = []string{"Variable_name", "Value"}
	case ast.ShowCollation:
		names = []string{"Collation", "Charset", "Id", "Default", "Compiled", "Sortlen"}
	case ast.ShowCreateTable, ast.ShowCreateSequence:
		names = []string{"Table", "Create Table"}
	case ast.ShowCreatePlacementPolicy:
		names = []string{"Policy", "Create Policy"}
	case ast.ShowCreateResourceGroup:
		names = []string{"Resource_Group", "Create Resource Group"}
	case ast.ShowCreateUser:
		names = []string{fmt.Sprintf("CREATE USER for %s", show.User)}
	case ast.ShowCreateView:
		names = []string{"View", "Create View", "character_set_client", "collation_connection"}
	case ast.ShowCreateDatabase:
		names = []string{"Database", "Create Database"}
	// unsupported on cloud
	// case ast.ShowDrainerStatus:
	// 	names = []string{"NodeID", "Address", "State", "Max_Commit_Ts", "Update_Time"}
	case ast.ShowGrants:
		names = []string{"Grants for User"}
	case ast.ShowIndex:
		names = []string{"Table", "Non_unique", "Key_name", "Seq_in_index",
			"Column_name", "Collation", "Cardinality", "Sub_part", "Packed",
			"Null", "Index_type", "Comment", "Index_comment", "Visible", "Expression", "Clustered"}
	case ast.ShowPlugins:
		names = []string{"Name", "Status", "Type", "Library", "License", "Version"}
	case ast.ShowProcessList:
		names = []string{"Id", "User", "Host", "db", "Command", "Time", "State", "Info"}
	// case ast.ShowPumpStatus:
	// 	names = []string{"NodeID", "Address", "State", "Max_Commit_Ts", "Update_Time"}
	case ast.ShowStatsMeta:
		names = []string{"Db_name", "Table_name", "Partition_name", "Update_time", "Modify_count", "Row_count"}
	case ast.ShowStatsExtended:
		names = []string{"Db_name", "Table_name", "Stats_name", "Column_names", "Stats_type", "Stats_val", "Last_update_version"}
	case ast.ShowStatsHistograms:
		names = []string{"Db_name", "Table_name", "Partition_name", "Column_name", "Is_index", "Update_time", "Distinct_count", "Null_count", "Avg_col_size", "Correlation",
			"Load_status", "Total_mem_usage", "Hist_mem_usage", "Topn_mem_usage", "Cms_mem_usage"}
	case ast.ShowStatsBuckets:
		names = []string{"Db_name", "Table_name", "Partition_name", "Column_name", "Is_index", "Bucket_id", "Count",
			"Repeats", "Lower_Bound", "Upper_Bound", "Ndv"}
	case ast.ShowStatsTopN:
		names = []string{"Db_name", "Table_name", "Partition_name", "Column_name", "Is_index", "Value", "Count"}
	case ast.ShowStatsHealthy:
		names = []string{"Db_name", "Table_name", "Partition_name", "Healthy"}
	case ast.ShowHistogramsInFlight:
		names = []string{"HistogramsInFlight"}
	case ast.ShowColumnStatsUsage:
		names = []string{"Db_name", "Table_name", "Partition_name", "Column_name", "Last_used_at", "Last_analyzed_at"}
	case ast.ShowStatsLocked:
		names = []string{"Db_name", "Table_name", "Partition_name", "Status"}
	case ast.ShowProfiles:
		names = []string{"Query_ID", "Duration", "Query"}
	case ast.ShowMasterStatus:
		names = []string{"File", "Position", "Binlog_Do_DB", "Binlog_Ignore_DB", "Executed_Gtid_Set"}
	case ast.ShowPrivileges:
		names = []string{"Privilege", "Context", "Comment"}
	case ast.ShowBindings:
		names = []string{"Original_sql", "Bind_sql", "Default_db", "Status", "Create_time", "Update_time", "Charset", "Collation", "Source", "Sql_digest", "Plan_digest"}
	case ast.ShowBindingCacheStatus:
		names = []string{"bindings_in_cache", "bindings_in_table", "memory_usage", "memory_quota"}
	case ast.ShowAnalyzeStatus:
		names = []string{"Table_schema", "Table_name", "Partition_name", "Job_info", "Processed_rows", "Start_time",
			"End_time", "State", "Fail_reason", "Instance", "Process_ID", "Remaining_seconds", "Progress", "Estimated_total_rows"}
	case ast.ShowBuiltins:
		names = []string{"Supported_builtin_functions"}
	case ast.ShowBackups, ast.ShowRestores:
		names = []string{"Id", "Destination", "State", "Progress", "Queue_time", "Execution_time", "Finish_time", "Connection", "Message"}
	case ast.ShowPlacementLabels:
		names = []string{"Key", "Values"}
	case ast.ShowPlacement, ast.ShowPlacementForDatabase, ast.ShowPlacementForTable, ast.ShowPlacementForPartition:
		names = []string{"Target", "Placement", "Scheduling_State"}
	case ast.ShowSessionStates:
		names = []string{"Session_states", "Session_token"}
	case ast.ShowImportJobs:
		names = []string{"Job_ID", "Create_Time", "Start_Time", "End_Time",
			"Data_Source", "Target_Table", "Created_By",
			"Status", "Source_File_Size", "Imported_Rows",
			"Result_Message"}
	}
	result := make(Columns, 0, 4)
	for _, name := range names {
		result = append(result, &Column{ColumnName: name})
	}
	return result, nil
}