package dbhubsdk

import (
	"context"
	"database/sql"
	"dataservices/commons/constants"
	"dataservices/commons/contextvar"
	"dataservices/commons/database"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/internal/entity"
	"encoding/hex"
	"fmt"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
	"unsafe"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
	mysql "github.com/tidbcloud/go-mysql-driver"
)

var ErrInvalidConn = mysql.ErrInvalidConn

type StmtType string

const (
	SELECT          StmtType = "select"
	INSERT          StmtType = "insert"
	UPDATE          StmtType = "update"
	DELETE          StmtType = "delete"
	DROPDATABASE    StmtType = "drop_database"
	DROPTABLE       StmtType = "drop_table"
	DROPINDEX       StmtType = "drop_index"
	ALTERTABLE      StmtType = "alter_table"
	SHOW            StmtType = "show"
	EXPLAIN         StmtType = "explain"
	USE             StmtType = "use"
	SET             StmtType = "set"
	CREATETABLE     StmtType = "create_table"
	CREATETEMPTABLE StmtType = "create_temporary_table"
	LOCKTABLE       StmtType = "lock_table"
	BEGIN           StmtType = "begin"
	COMMIT          StmtType = "commit"
	ROLLBACK        StmtType = "rollback"
	OTHER           StmtType = "other"
)

func GetStmtTypeBySql(ctx context.Context, sql string) (StmtType, error) {
	var (
		stmtType StmtType
		err      error
	)

	stmtNode, err := parser.New().ParseOneStmt(sql, "", "")
	if err != nil {
		err = errs.NewBadRequest("Parse stmt failed, "+err.Error(), nil)
		return stmtType, err
	}
	return GetStmtType(ctx, stmtNode), nil
}

func GetStmtType(ctx context.Context, stmtNode ast.StmtNode) StmtType {

	if stmtNode == nil {
		return OTHER
	}

	switch node := stmtNode.(type) { //nolint:gosimple // ignore
	case *ast.InsertStmt:
		return INSERT
	case *ast.SelectStmt, *ast.SetOprStmt:
		return SELECT
	case *ast.UpdateStmt:
		return UPDATE
	case *ast.DeleteStmt:
		return DELETE
	case *ast.ShowStmt:
		return SHOW
	case *ast.ExplainStmt:
		return EXPLAIN
	case *ast.DropDatabaseStmt:
		return DROPDATABASE
	case *ast.DropTableStmt:
		return DROPTABLE
	case *ast.DropIndexStmt:
		return DROPINDEX
	case *ast.AlterTableStmt:
		return ALTERTABLE
	case *ast.UseStmt:
		return USE
	case *ast.SetStmt:
		return SET
	case *ast.CreateTableStmt:
		if node.TemporaryKeyword == ast.TemporaryNone {
			return CREATETABLE
		} else {
			return CREATETEMPTABLE
		}
	case *ast.LockTablesStmt:
		return LOCKTABLE
	case *ast.BeginStmt:
		return BEGIN
	case *ast.CommitStmt:
		return COMMIT
	case *ast.RollbackStmt:
		return ROLLBACK
	case *ast.NonTransactionalDMLStmt:
		if node.DryRun != ast.NoDryRun {
			return SELECT
		}
		return GetStmtType(ctx, node.DMLStmt)
	}
	tilog.Info(ctx, fmt.Sprintf("Other StmtNode: %T", stmtNode))
	return OTHER
}

func IsQueryStmt(ctx context.Context, stmtNode ast.StmtNode) bool {

	if stmtNode == nil {
		return false
	}

	switch node := stmtNode.(type) { //nolint:gosimple // ignore
	case *ast.AdminStmt:
		switch node.Tp {
		case ast.AdminShowDDL,
			ast.AdminShowDDLJobQueries,
			ast.AdminShowDDLJobQueriesWithRange,
			ast.AdminShowDDLJobs,
			ast.AdminShowNextRowID,
			ast.AdminShowSlow,
			ast.AdminCheckIndex,
			ast.AdminCheckIndexRange,
			ast.AdminCheckTable,
			ast.AdminChecksumTable,
			ast.AdminCleanupIndex,
			ast.AdminPauseDDLJobs,
			ast.AdminRecoverIndex,
			ast.AdminResumeDDLJobs,
			ast.AdminCancelDDLJobs:
			return true
		}
	case
		*ast.CalibrateResourceStmt,
		*ast.ExecuteStmt,
		*ast.ExplainForStmt,
		*ast.ExplainStmt,
		*ast.SelectStmt,
		*ast.ShowStmt,
		*ast.SplitRegionStmt,
		*ast.TraceStmt,
		*ast.SetOprStmt,
		*ast.PlanReplayerStmt:
		return true
	case *ast.NonTransactionalDMLStmt:
		return node.DryRun != ast.NoDryRun
	}
	return false
}

func GetDatabases(ctx context.Context, stmtNode ast.StmtNode) []string {
	if stmtNode == nil {
		return nil
	}
	databases := getSchema(stmtNode)
	exists := make(map[string]struct{})
	result := make([]string, 0, len(databases))
	for _, database := range databases {
		if _, ok := exists[database]; ok {
			continue
		}
		exists[database] = struct{}{}
		result = append(result, database)
	}
	return result
}

func getSchema(node ast.Node) []string {
	if node == nil {
		return nil
	}
	var ans []string
	switch astNode := node.(type) {

	case *ast.SelectStmt:
		if astNode.From == nil {
			return nil
		}
		ans = append(ans, getSchema(astNode.From.TableRefs)...)
	case *ast.SetOprStmt:
		if astNode.SelectList == nil || len(astNode.SelectList.Selects) == 0 {
			return nil
		}
		for _, selectNode := range astNode.SelectList.Selects {
			ans = append(ans, getSchema(selectNode)...)
		}
	case *ast.SetOprSelectList:
		if len(astNode.Selects) == 0 {
			return nil
		}
		for _, selectNode := range astNode.Selects {
			ans = append(ans, getSchema(selectNode)...)
		}

	case *ast.Join:
		ans = append(ans, getSchema(astNode.Left)...)
		ans = append(ans, getSchema(astNode.Right)...)
	case *ast.TableSource:
		ans = append(ans, getSchema(astNode.Source)...)
	case *ast.TableName:
		ans = append(ans, astNode.Schema.O)
	}
	return ans
}

func executeSQL(ctx context.Context, db *sql.Conn, sqlStr string, sqlType StmtType, isQueryStmt bool) (rows *sql.Rows, lastID int64, affect int64, err error) {
	startTime := time.Now()
	logger := contextvar.Logger(ctx)
	defer func() {
		cost := time.Since(startTime)
		logger.Info("executeSQL completed", zap.Time("startTime", startTime), zap.String("sql", sqlStr), zap.String("cost", cost.String()), zap.String("sqlType", string(sqlType)), zap.Bool("isQueryStmt", isQueryStmt), zap.Error(err))
	}()

	if isQueryStmt {
		rows, err = db.QueryContext(ctx, sqlStr)
		return
	}

	var result sql.Result
	result, err = db.ExecContext(ctx, sqlStr)
	if err != nil {
		return
	}
	affect, err = result.RowsAffected()
	if err != nil {
		return
	}
	lastID, err = result.LastInsertId()
	if err != nil {
		return
	}
	return
}

func queryDBByPrepare(ctx context.Context, db *sql.Conn, sql string, args []interface{}, limits uint64) (res *QueryRes, err error) {
	startTime := time.Now()
	logger := contextvar.Logger(ctx)
	defer func() {
		var (
			cost    = time.Since(startTime)
			dumpRes QueryRes
		)
		if res != nil {
			dumpRes = *res
			dumpRes.Columns = nil
			dumpRes.Rows = nil
		}
		logger.Info("queryDBByPrepare completed",
			zap.Time("startTime", startTime),
			zap.Duration("cost", cost),
			zap.String("sql", sql),
			zap.Any("dumpRes", dumpRes),
			zap.Error(err))
	}()
	stmtNode, err := parser.New().ParseOneStmt(sql, "", "")
	if err != nil {
		logger.Error("ParseOneStmt failed", zap.Error(err), zap.String("sql", sql))
	}
	sqlType := GetStmtType(ctx, stmtNode)
	isQueryStmt := IsQueryStmt(ctx, stmtNode)
	databases := GetDatabases(ctx, stmtNode)
	startMs := time.Now().UnixNano() / int64(time.Millisecond)
	rows, lastInsertId, affect, err := executeSQLByPrepare(ctx, db, sql, args, sqlType, isQueryStmt)
	endMs := time.Now().UnixNano() / int64(time.Millisecond)
	res = &QueryRes{
		StartMs:      startMs,
		EndMs:        endMs,
		Affect:       affect,
		Databases:    databases,
		StmtType:     sqlType,
		LastInsertID: lastInsertId,
	}
	if stmtNode != nil {
		if useStmt, ok := stmtNode.(*ast.UseStmt); ok && useStmt != nil {
			res.UseDatabase = useStmt.DBName
		} else if dropStmt, ok := stmtNode.(*ast.DropDatabaseStmt); ok && dropStmt != nil {
			res.DropDatabase = dropStmt.Name.String()
		}
	}
	if err != nil {
		if errMySQL, ok := err.(*mysql.MySQLError); ok {
			res.Code = errMySQL.Number
			res.Message = errMySQL.Message
			return res, nil
		}
		return nil, err
	}

	if rows != nil {
		ScanRows(ctx, rows, limits, res)
		rows.Close()
	}
	return res, nil
}

func executeSQLByPrepare(ctx context.Context, db *sql.Conn, sql string, args []interface{}, sqlType StmtType, isQueryStmt bool) (rows *sql.Rows, lastInsertId, affect int64, err error) {
	var (
		startTime = time.Now()
		logger    = contextvar.Logger(ctx)
		isPrepare = len(args) > 0
	)
	logger.Info("execute sql begin", zap.String("sql", sql), zap.Any("args", args),
		zap.String("sqlType", string(sqlType)), zap.Bool("isQueryStmt", isQueryStmt),
		zap.Bool("isPrepare", isPrepare))
	defer func() {
		logger.Info("execute sql completed", zap.Time("startTime", startTime), zap.String("sql", sql),
			zap.Duration("cost", time.Since(startTime)), zap.String("sqlType", string(sqlType)), zap.Bool("isQueryStmt", isQueryStmt),
			zap.Bool("isPrepare", isPrepare), zap.Error(err))
	}()

	if len(args) == 0 {
		return executeSQL(ctx, db, sql, sqlType, isQueryStmt)
	}

	stmt, err := db.PrepareContext(ctx, sql)
	metrics.PrepareStmtDuration.WithLabelValues(strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
	if err != nil {
		return
	}

	startTime = time.Now()
	if isQueryStmt {
		rows, err = stmt.QueryContext(ctx, args...)
		metrics.StmtExecuteDuration.WithLabelValues(strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
		return
	}
	result, err := stmt.ExecContext(ctx, args...)
	metrics.StmtExecuteDuration.WithLabelValues(strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
	if err != nil {
		return
	}
	affect, err = result.RowsAffected()
	if err != nil {
		return
	}
	lastInsertId, err = result.LastInsertId()
	return
}

// convertCellValue converts a value scanned from a sql.Rows into its string
// representation, dispatching on the Go type and the column's database type
// name. Drivers return MySQL/TiDB BIT(N) values as raw big-endian bytes
// (e.g. b'0' -> []byte{0x00}, b'1' -> []byte{0x01}); rendering those bytes as
// text yields invisible control characters that appear blank in the UI, so
// BIT is decoded as an unsigned integer instead.
func convertCellValue(v interface{}, dbTypeName string) string {
	if v == nil {
		return ""
	}
	switch val := v.(type) {
	case time.Time:
		return val.Format("2006-01-02 15:04:05")
	case []byte:
		switch dbTypeName {
		case "BIT":
			var num uint64
			for _, b := range val {
				num = (num << 8) | uint64(b)
			}
			return strconv.FormatUint(num, 10)
		case "BINARY", "VARBINARY", "TINYBLOB", "BLOB", "MEDIUMBLOB", "LONGBLOB", "VARCHAR":
			if utf8.Valid(val) {
				return Byte2Str(val)
			}
			return strings.ToUpper("0x" + hex.EncodeToString(val))
		default:
			return Byte2Str(val)
		}
	default:
		return fmt.Sprintf("%v", val)
	}
}

func ScanRows(ctx context.Context, rows *sql.Rows, limits uint64, result *QueryRes) {
	if rows == nil {
		return
	}
	startTime := time.Now()
	logger := contextvar.Logger(ctx)
	defer func() {
		cost := time.Since(startTime)
		logger.Info("ScanRows completed", zap.Time("startTime", startTime), zap.String("cost", cost.String()))
		metrics.ScanRowsDuration.Observe(cost.Seconds())
	}()
	rowNum := int64(0)

	array := make([][]string, 0)
	colTypes, _ := rows.ColumnTypes()
	var rowParam = make([]interface{}, len(colTypes))
	var rowValue = make([]interface{}, len(colTypes))

	for i, colType := range colTypes {
		rowValue[i] = reflect.New(colType.ScanType())
		rowParam[i] = reflect.ValueOf(&rowValue[i]).Interface()

	}
	var columns []*Column

	//columns
	for _, col := range colTypes {
		nullable, _ := col.Nullable()
		dataType := col.DatabaseTypeName()
		isTimeType := false
		timeType := ""
		if dataType == "DATETIME" {
			isTimeType = true
			timeType = "day|week|month|year"
		}
		param := &Column{
			Col:        col.Name(),
			DataType:   col.DatabaseTypeName(),
			Nullable:   nullable,
			IsTimeType: isTimeType,
			TimeType:   timeType,
		}
		columns = append(columns, param)
	}

	// rows
	for rows.Next() {
		rowNum++
		if limits > 0 && rowNum > int64(limits) {
			continue
		}
		rows.Scan(rowParam...)
		row := make([]string, 0)
		for i, ct := range colTypes {
			row = append(row, convertCellValue(rowValue[i], ct.DatabaseTypeName()))
		}
		array = append(array, row)

	}
	result.Columns = columns
	result.Rows = array
	result.RowNum = rowNum
	return
}

// Byte2Str []byte to string
func Byte2Str(b []byte) string {
	return *(*string)(unsafe.Pointer(&b))
}

func connectDB(ctx context.Context, user *TidbUserMeta, in *entity.ConnectionInput, useTls bool) (*sql.DB, error) {
	var err error
	startTime := time.Now()
	logger := contextvar.Logger(ctx)
	defer func() {
		cost := time.Since(startTime)
		logger.Info("try to connectDB process completed", zap.Uint64("cluster", user.ClusterId), zap.Time("startTime", startTime), zap.String("cost", cost.String()), zap.Error(err))
		metrics.DatabaseConnDuration.WithLabelValues(strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
	}()

	connectionSource := "tcds"
	if in.ConnectionSource != "" {
		connectionSource = connectionSource + "-" + in.ConnectionSource
	}

	database := in.Database
	maxIdleTime := in.ExpireTime
	maxOpenConns, maxIdleConns := 1, 1

	if maxIdleTime == 0 {
		maxIdleTime = time.Minute
	}
	if in.MaxIdleConns != nil {
		maxIdleConns = *in.MaxIdleConns
	}
	if in.MaxOpenConns != nil {
		maxOpenConns = *in.MaxOpenConns
	}

	logger.Info(fmt.Sprintf("BuildConnection:%s:psw@tcp(%s:%d)/%s?timeout=30s&readTimeout=60s&writeTimeout=60s&allowCleartextPasswords=true&connectionAttributes=program_name:%s",
		user.ServerlessUser,
		user.Host,
		user.Port,
		database,
		connectionSource,
	), zap.Bool("tls", useTls))
	// Params
	url := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?timeout=30s&readTimeout=60s&writeTimeout=60s&allowCleartextPasswords=true&connectionAttributes=program_name:%s",
		user.ServerlessUser,
		user.ServerlessPasswd,
		user.Host,
		user.Port,
		database,
		connectionSource,
	)
	if useTls {
		url += "&tls=true"
	}
	// Connect
	sqlDB, err := sql.Open("go-mysql", url)
	if err != nil {
		return nil, err
	}

	sqlDB.SetMaxIdleConns(maxIdleConns)
	sqlDB.SetMaxOpenConns(maxOpenConns)
	sqlDB.SetConnMaxIdleTime(maxIdleTime)
	return sqlDB, nil
}

var (
	sqlArgsRegexp     = regexp.MustCompile(constants.SqlArgsRegexp)
	sqlCommentsRegexp = regexp.MustCompile(constants.SqlCommentsRegexp)
)

func RWMode(ctx context.Context, sql string) string {
	sql = sqlCommentsRegexp.ReplaceAllString(sql, "")
	sql = sqlArgsRegexp.ReplaceAllString(sql, "?")
	if sql == "" {
		return constants.SqlModeINVALID
	}
	stmts, _, err := parser.New().ParseSQL(sql)
	if err != nil {
		logger.Warn(ctx, "getSqlRwMode failed, invalid sql", zap.Error(err))
		return constants.SqlModeINVALID
	}
	isReadOnly := true
	for _, stmt := range stmts {
		isReadOnly = isReadOnly && isStmtReadOnly(stmt)
	}
	var mode string
	if isReadOnly {
		mode = constants.SqlModeReadOnly
	} else {
		mode = constants.SqlModeReadWrite
	}
	return mode
}

func isStmtReadOnly(node ast.StmtNode) bool {
	if ast.IsReadOnly(node, true) {
		return true
	}
	_, ok := node.(*ast.UseStmt)
	return ok
}

func IsCausedByUser(code uint16, message string) bool {
	return database.IsCausedByUser(code, message)
}