// Copyright 2018 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 main

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/pingcap/enterprise-toolkit/o2t-sync-diff/pkg/dbutil"
	"github.com/pingcap/enterprise-toolkit/o2t-sync-diff/pkg/diff"
	router "github.com/pingcap/enterprise-toolkit/o2t-sync-diff/pkg/table-router"
	"github.com/pingcap/enterprise-toolkit/o2t-sync-diff/pkg/utils"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	tidbconfig "github.com/pingcap/tidb/config"
	"go.uber.org/zap"
)

var fixsqlWarning = "请注意！！！本程序会将空字符串 '' 与 NULL 识别为不一致的数据并进行修复，当被修复库是 TiDB 时，请再三和应用方确认是否将 TiDB 中的空字符串 '' 修复为 NULL，若无需修复，请找到这些修复语句并删除它们。（搜索关键字 NULL 确认，请注意大小写）\n请人工确认修复sql的正确性"

// Diff contains two sql DB, used for comparing.
type Diff struct {
	sourceDBs         map[string]DBConfig
	targetDB          DBConfig
	chunkSize         int
	sample            int
	checkThreadCount  int
	useChecksum       bool
	useCheckpoint     bool
	onlyUseChecksum   bool
	ignoreDataCheck   bool
	ignoreStructCheck bool
	ignoreStats       bool
	tables            map[string]map[string]*TableConfig
	fixSQLFile        *os.File
	haveOutputWarning bool
	fixTarget         string
	EmptyStringNullCompareSwitch bool
	Ascii0Switch 	  bool

	report         *Report
	tidbInstanceID string
	tableRouter    *router.Table
	cpDB           *sql.DB

	// DM's subtask config
	//subTaskCfgs []*config.SubTaskConfig

	ctx context.Context

	OracleSchema string
	OracleSCN    string
}

// NewDiff returns a Diff instance.
func NewDiff(ctx context.Context, cfg *Config) (diff *Diff, err error) {
	diff = &Diff{
		sourceDBs:         make(map[string]DBConfig),
		chunkSize:         cfg.ChunkSize,
		sample:            cfg.Sample,
		checkThreadCount:  cfg.CheckThreadCount,
		useChecksum:       cfg.UseChecksum,
		useCheckpoint:     cfg.UseCheckpoint,
		onlyUseChecksum:   cfg.OnlyUseChecksum,
		ignoreDataCheck:   cfg.IgnoreDataCheck,
		ignoreStructCheck: cfg.IgnoreStructCheck,
		ignoreStats:       cfg.IgnoreStats,
		tables:            make(map[string]map[string]*TableConfig),
		report:            NewReport(),
		ctx:               ctx,
		OracleSchema:      cfg.Tables[0].OracleSchema,
		fixTarget:         cfg.FixTarget,
		EmptyStringNullCompareSwitch: cfg.EmptyStringNullCompareSwitch,
		Ascii0Switch: 	   cfg.Ascii0Switch,
		OracleSCN:         cfg.SourceDBCfg[0].OracleSCN,
		haveOutputWarning: false,
	}

	if err = diff.init(cfg); err != nil {
		diff.Close()
		return nil, errors.Trace(err)
	}

	return diff, nil
}

func (df *Diff) init(cfg *Config) (err error) {
	setTiDBCfg()

	// create connection for source.
	if err = df.CreateDBConn(cfg); err != nil {
		return errors.Trace(err)
	}

	if err = df.AdjustTableConfig(cfg); err != nil {
		return errors.Trace(err)
	}

	df.fixSQLFile, err = os.Create(cfg.FixSQLFile)
	if err != nil {
		return errors.Trace(err)
	}

	return nil
}

// CreateDBConn creates db connections for source and target.
func (df *Diff) CreateDBConn(cfg *Config) (err error) {
	// create connection for target.
	cfg.TargetDBCfg.Conn, err = diff.CreateDB(df.ctx, cfg.TargetDBCfg.DBConfig, nil, cfg.CheckThreadCount*2)
	if err != nil {
		return errors.Errorf("create target db %s error %v", cfg.TargetDBCfg.DBConfig.String(), err)
	}
	df.targetDB = cfg.TargetDBCfg

	targetTZOffset, err := dbutil.GetTimeZoneOffset(df.ctx, cfg.TargetDBCfg.Conn)
	if err != nil {
		return errors.Annotatef(err, "fetch target db %s time zone offset failed", cfg.TargetDBCfg.DBConfig.String())
	}
	vars := map[string]string{
		"time_zone": dbutil.FormatTimeZoneOffset(targetTZOffset),
	}

	for _, source := range cfg.SourceDBCfg {
		// connect source db with target db time_zone
		source.Conn, err = diff.CreateDB(df.ctx, source.DBConfig, vars, cfg.CheckThreadCount*2)
		if err != nil {
			return errors.Annotatef(err, "create source db %s failed", source.DBConfig.String())
		}
		df.sourceDBs[source.InstanceID] = source
	}

	df.cpDB, err = diff.CreateDBForCP(df.ctx, cfg.TargetDBCfg.DBConfig)
	if err != nil {
		return errors.Errorf("create checkpoint db %s error %v", cfg.TargetDBCfg.DBConfig.String(), err)
	}

	return nil
}

// AdjustTableConfig adjusts the table's config by check-tables and table-config.
func (df *Diff) AdjustTableConfig(cfg *Config) (err error) {

	df.tableRouter, err = router.NewTableRouter(false, cfg.TableRules)
	if err != nil {
		return errors.Trace(err)
	}

	allTablesMap, err := df.GetAllTables(cfg)
	if err != nil {
		return errors.Trace(err)
	}

	// get all source table's matched target table
	// target database name => target table name => all matched source table instance
	sourceTablesMap := make(map[string]map[string][]TableInstance)
	for instanceID, allSchemas := range allTablesMap {
		if instanceID == df.targetDB.InstanceID {
			continue
		}

		for schema, allTables := range allSchemas {
			for table := range allTables {
				targetSchema, targetTable, err := df.tableRouter.Route(schema, table)
				if err != nil {
					return errors.Errorf("get route result for %s.%s.%s failed, error %v", instanceID, schema, table, err)
				}

				if _, ok := sourceTablesMap[targetSchema]; !ok {
					sourceTablesMap[targetSchema] = make(map[string][]TableInstance)
				}

				if _, ok := sourceTablesMap[targetSchema][targetTable]; !ok {
					sourceTablesMap[targetSchema][targetTable] = make([]TableInstance, 0, 1)
				}
				mapSchema := schema
				if df.OracleSchema != "" {
					mapSchema = df.OracleSchema
				}
				sourceTablesMap[targetSchema][targetTable] = append(sourceTablesMap[targetSchema][targetTable], TableInstance{
					InstanceID: instanceID,
					Schema:     mapSchema,
					Table:      table,
				})
			}
		}
	}

	// fill the table information.
	// will add default source information, don't worry, we will use table config's info replace this later.
	for _, schemaTables := range cfg.Tables {
		df.tables[schemaTables.Schema] = make(map[string]*TableConfig)
		tables := make([]string, 0, len(schemaTables.Tables))
		allTables, ok := allTablesMap[df.targetDB.InstanceID][schemaTables.Schema]
		if !ok {
			return errors.NotFoundf("schema %s.%s", df.targetDB.InstanceID, schemaTables.Schema)
		}

		for _, table := range schemaTables.Tables {
			matchedTables, err := df.GetMatchTable(df.targetDB, schemaTables.Schema, table, allTables)
			if err != nil {
				return errors.Trace(err)
			}

			//exclude those in "exclude-tables"
			for _, t := range matchedTables {
				if df.InExcludeTables(schemaTables.ExcludeTables, t) {
					continue
				} else {
					tables = append(tables, t)
				}
			}
		}

		for _, tableName := range tables {
			tableInfo, err := dbutil.GetTableInfo(df.ctx, df.targetDB.Conn, schemaTables.Schema, tableName)
			if err != nil {
				return errors.Errorf("get table %s.%s's information error %s", schemaTables.Schema, tableName, errors.ErrorStack(err))
			}

			if _, ok := df.tables[schemaTables.Schema][tableName]; ok {
				log.Error("duplicate config for one table", zap.String("table", dbutil.TableName(schemaTables.Schema, tableName)))
				continue
			}

			sourceTables := make([]TableInstance, 0, 1)
			if _, ok := sourceTablesMap[schemaTables.Schema][tableName]; ok {
				log.Info("find matched source tables", zap.Reflect("source tables", sourceTablesMap[schemaTables.Schema][tableName]), zap.String("target schema", schemaTables.Schema), zap.String("table", tableName))
				sourceTables = sourceTablesMap[schemaTables.Schema][tableName]
			} else {
				// use same database name and table name if OracleSchema not setup
				mapSchema := schemaTables.Schema
				if df.OracleSchema != "" {
					mapSchema = df.OracleSchema
				}
				sourceTables = append(sourceTables, TableInstance{
					InstanceID: cfg.SourceDBCfg[0].InstanceID,
					Schema:     mapSchema,
					Table:      tableName,
				})
			}

			df.tables[schemaTables.Schema][tableName] = &TableConfig{
				TableInstance: TableInstance{
					Schema: schemaTables.Schema,
					Table:  tableName,
				},
				TargetTableInfo: tableInfo,
				Range:           "1=1",
				OracleRange:     "1=1",
				SourceTables:    sourceTables,
			}
		}
	}

	for _, table := range cfg.TableCfgs {
		if _, ok := df.tables[table.Schema]; !ok {
			return errors.NotFoundf("schema %s in check tables", table.Schema)
		}
		if _, ok := df.tables[table.Schema][table.Table]; !ok {
			return errors.NotFoundf("table %s.%s in check tables", table.Schema, table.Table)
		}

		sourceTables := make([]TableInstance, 0, len(table.SourceTables))
		for _, sourceTable := range table.SourceTables {
			if _, ok := df.sourceDBs[sourceTable.InstanceID]; !ok {
				return errors.Errorf("unkonwn database instance id %s", sourceTable.InstanceID)
			}

			sourceTables = append(sourceTables, TableInstance{
				InstanceID: sourceTable.InstanceID,
				Schema:     sourceTable.Schema,
				Table:      sourceTable.Table,
			})
			log.Info("modify matched source tables", zap.Reflect("source tables", sourceTables), zap.String("target schema", table.Schema), zap.String("table", table.Table))
		}

		if len(sourceTables) != 0 {
			df.tables[table.Schema][table.Table].SourceTables = sourceTables
		}
		if table.Range != "" {
			df.tables[table.Schema][table.Table].Range = table.Range
		}
		if table.OracleRange != "" {
			df.tables[table.Schema][table.Table].OracleRange = table.OracleRange
		}
		df.tables[table.Schema][table.Table].OracleHint = table.OracleHint
		if table.TidbHint != "" {
			df.tables[table.Schema][table.Table].TidbHint = table.TidbHint
		}
		df.tables[table.Schema][table.Table].IgnoreColumns = table.IgnoreColumns
		df.tables[table.Schema][table.Table].OnlyCheckColumns = table.OnlyCheckColumns
		df.tables[table.Schema][table.Table].Fields = table.Fields
		df.tables[table.Schema][table.Table].Collation = table.Collation
	}

	// we need to increase max open connections for upstream, because one chunk needs accessing N shard tables in one
	// upstream, and there are `CheckThreadCount` processing chunks. At most we need N*`CheckThreadCount` connections
	// for an upstream
	// instanceID -> max number of upstream shard tables every target table
	maxNumShardTablesOneRun := map[string]int{}
	for _, targetTables := range df.tables {
		for _, sourceCfg := range targetTables {
			upstreamCount := map[string]int{}
			for _, sourceTables := range sourceCfg.SourceTables {
				upstreamCount[sourceTables.InstanceID]++
			}
			for id, count := range upstreamCount {
				if count > maxNumShardTablesOneRun[id] {
					maxNumShardTablesOneRun[id] = count
				}
			}
		}
	}

	for instanceId, count := range maxNumShardTablesOneRun {
		db := df.sourceDBs[instanceId].Conn
		if db == nil {
			return errors.Errorf("didn't found sourceDB for instance %s", instanceId)
		}
		log.Info("will increase connection configurations for DB of instance",
			zap.String("instance id", instanceId),
			zap.Int("connection limit", count*df.checkThreadCount))
		db.SetMaxOpenConns(count * df.checkThreadCount * 2)
		db.SetMaxIdleConns(count * df.checkThreadCount * 2)
	}

	return nil
}

// GetAllTables get all tables in all databases.
func (df *Diff) GetAllTables(cfg *Config) (map[string]map[string]map[string]interface{}, error) {
	// instanceID => schema => table
	allTablesMap := make(map[string]map[string]map[string]interface{})

	allTablesMap[df.targetDB.InstanceID] = make(map[string]map[string]interface{})
	targetSchemas, err := dbutil.GetSchemas(df.ctx, df.targetDB.Conn)
	if err != nil {
		return nil, errors.Annotatef(err, "get schemas from %s", df.targetDB.InstanceID)
	}
	for _, schema := range targetSchemas {
		allTables, err := dbutil.GetTables(df.ctx, df.targetDB.Conn, schema)
		if err != nil {
			return nil, errors.Annotatef(err, "get tables from %s.%s", df.targetDB.InstanceID, schema)
		}
		allTablesMap[df.targetDB.InstanceID][schema] = utils.SliceToMap(allTables)
	}

	//for Oracle db source, it do not support table-rules.
	//it have no multi sources; so one target schema correspond to  one source schema
	if cfg.SourceDBCfg[0].Type == dbutil.TypeOracle {
		for _, source := range df.sourceDBs {
			allTablesMap[source.InstanceID] = make(map[string]map[string]interface{})
			for _, schema := range targetSchemas {
				allTablesMap[source.InstanceID][schema] = allTablesMap[df.targetDB.InstanceID][schema]
			}
		}
		return allTablesMap, nil
	}

	//for TiDB
	for _, source := range df.sourceDBs {
		allTablesMap[source.InstanceID] = make(map[string]map[string]interface{})
		sourceSchemas, err := dbutil.GetSchemas(df.ctx, source.Conn)
		if err != nil {
			return nil, errors.Annotatef(err, "get schemas from %s", source.InstanceID)
		}

		for _, schema := range sourceSchemas {
			allTables, err := dbutil.GetTables(df.ctx, source.Conn, schema)
			if err != nil {
				return nil, errors.Annotatef(err, "get tables from %s.%s", source.InstanceID, schema)
			}
			allTablesMap[source.InstanceID][schema] = utils.SliceToMap(allTables)
		}
	}

	return allTablesMap, nil
}

// GetMatchTable returns all the matched table.
func (df *Diff) GetMatchTable(db DBConfig, schema, table string, allTables map[string]interface{}) ([]string, error) {
	tableNames := make([]string, 0, 1)

	if table[0] == '~' {
		tableRegex := regexp.MustCompile(fmt.Sprintf("(?i)%s", table[1:]))
		for tableName := range allTables {
			if !tableRegex.MatchString(tableName) {
				continue
			}
			tableNames = append(tableNames, tableName)
		}
	} else {
		if _, ok := allTables[table]; ok {
			tableNames = append(tableNames, table)
		} else {
			return nil, errors.Errorf("%s.%s not found in %s", schema, table, db.InstanceID)
		}
	}

	return tableNames, nil
}

// Close closes file and database connection.
func (df *Diff) Close() {
	if df.fixSQLFile != nil {
		df.fixSQLFile.Close()
	}

	for _, db := range df.sourceDBs {
		if db.Conn != nil {
			db.Conn.Close()
		}
	}

	if df.targetDB.Conn != nil {
		df.targetDB.Conn.Close()
	}

	if df.cpDB != nil {
		df.cpDB.Close()
	}
}

// Equal tests whether two database have same data and schema.
func (df *Diff) Equal() (err error) {
	defer df.Close()

	for _, schema := range df.tables {
		for _, table := range schema {
			sourceTables := make([]*diff.TableInstance, 0, len(table.SourceTables))
			for _, sourceTable := range table.SourceTables {
				sourceTableInstance := &diff.TableInstance{
					Conn:       df.sourceDBs[sourceTable.InstanceID].Conn,
					Schema:     sourceTable.Schema,
					Table:      sourceTable.Table,
					InstanceID: sourceTable.InstanceID,
					DBType:     df.sourceDBs[sourceTable.InstanceID].DBConfig.Type,
					Charset:    df.sourceDBs[sourceTable.InstanceID].DBConfig.Charset,
				}
				sourceTables = append(sourceTables, sourceTableInstance)
			}

			targetTableInstance := &diff.TableInstance{
				Conn:       df.targetDB.Conn,
				Schema:     table.Schema,
				Table:      table.Table,
				InstanceID: df.targetDB.InstanceID,
				DBType:     df.targetDB.DBConfig.Type,
				Charset:    df.targetDB.DBConfig.Charset,
			}

			// find tidb instance for getting statistical information to split chunk
			var tidbStatsSource *diff.TableInstance
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			if !df.ignoreStats {
				log.Info("use tidb stats to split chunks")
				isTiDB, err := dbutil.IsTiDB(ctx, targetTableInstance.Conn)
				if err != nil {
					log.Warn("judge instance is tidb failed", zap.Error(err))
				} else if isTiDB {
					tidbStatsSource = targetTableInstance
				} else if len(sourceTables) == 1 {
					isTiDB, err := dbutil.IsTiDB(ctx, sourceTables[0].Conn)
					if err != nil {
						log.Warn("judge instance is tidb failed", zap.Error(err))
					} else if isTiDB {
						tidbStatsSource = sourceTables[0]
					}
				}
			} else {
				log.Info("ignore tidb stats because of user setting")
			}

			td := &diff.TableDiff{
				SourceTables: sourceTables,
				TargetTable:  targetTableInstance,

				IgnoreColumns:    table.IgnoreColumns,
				OnlyCheckColumns: table.OnlyCheckColumns,

				Fields:            table.Fields,
				Range:             table.Range,
				OracleRange:       table.OracleRange,
				OracleHint:		   table.OracleHint,
				TidbHint:		   table.TidbHint,
				Collation:         table.Collation,
				ChunkSize:         df.chunkSize,
				Sample:            df.sample,
				CheckThreadCount:  df.checkThreadCount,
				UseChecksum:       df.useChecksum,
				UseCheckpoint:     df.useCheckpoint,
				OnlyUseChecksum:   df.onlyUseChecksum,
				IgnoreStructCheck: df.ignoreStructCheck,
				IgnoreDataCheck:   df.ignoreDataCheck,
				TiDBStatsSource:   tidbStatsSource,
				CpDB:              df.cpDB,
				FixTarget:         df.fixTarget,
				EmptyStringNullCompareSwitch: df.EmptyStringNullCompareSwitch,
				Ascii0Switch:      df.Ascii0Switch,
				OracleSCN:         df.OracleSCN,
			}
			log.Info("check target table", zap.Reflect("table ", targetTableInstance.Table),
				zap.Reflect("connection pool stat", sourceTables[0].Conn.Stats()))
			structEqual, dataEqual, err := td.Equal(df.ctx, func(dml string) error {
				var err error
				if df.haveOutputWarning {
					_, err = df.fixSQLFile.WriteString(fmt.Sprintf("%s\n", dml))
				} else {
					_, err = df.fixSQLFile.WriteString(fmt.Sprintf("%s\n%s\n", fixsqlWarning, dml))
					df.haveOutputWarning = true
				}
				return errors.Trace(err)
			})

			if err != nil {
				log.Error("check failed", zap.String("table", dbutil.TableName(table.Schema, table.Table)), zap.Error(err))
				df.report.SetTableMeetError(table.Schema, table.Table, err)
				df.report.FailedNum++
				continue
			}

			df.report.SetTableStructCheckResult(table.Schema, table.Table, structEqual)
			df.report.SetTableDataCheckResult(table.Schema, table.Table, dataEqual)
			if structEqual && dataEqual {
				df.report.PassNum++
			} else {
				df.report.FailedNum++
			}
		}
	}

	return
}

// Judge if a table is in "exclude-tables" list
func (df *Diff) InExcludeTables(exclude_tables []string, table string) bool {
	for _, exclude_table := range exclude_tables {
		if strings.EqualFold(exclude_table, table) {
			return true
		}
	}
	return false
}

func setTiDBCfg() {
	// to support long index key in TiDB
	tidbCfg := tidbconfig.GetGlobalConfig()
	// 3027 * 4 is the max value the MaxIndexLength can be set
	tidbCfg.MaxIndexLength = tidbconfig.DefMaxOfMaxIndexLength
	tidbCfg.Experimental.AllowsExpressionIndex = true
	tidbconfig.StoreGlobalConfig(tidbCfg)

	fmt.Println("set tidb cfg")
}