package import_svc

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/pkg/errno"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/clusterproxy"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/serverless-service/internal/domain/constant"
	"github.com/tidbcloud/serverless-service/internal/domain/entity"
	"github.com/tidbcloud/serverless-service/internal/domain/facade"
	"github.com/tidbcloud/serverless-service/internal/domain/repo"
	"github.com/tidbcloud/serverless-service/internal/errorcode"
	"github.com/tidbcloud/serverless-service/internal/utils"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
)

type ImportHelperService struct {
	importRepo                 repo.ImportRepo
	accountFacade              facade.AccountFacade
	encryptorFacade            facade.EncryptorFacade
	clusterProxyProviderFacade facade.ClusterProxyProviderFacade
}

func NewImportHelperService(
	importRepo repo.ImportRepo,
	accountFacade facade.AccountFacade,
	encryptorFacade facade.EncryptorFacade,
	clusterProxyProviderFacade facade.ClusterProxyProviderFacade,
) *ImportHelperService {
	return &ImportHelperService{
		importRepo:                 importRepo,
		accountFacade:              accountFacade,
		encryptorFacade:            encryptorFacade,
		clusterProxyProviderFacade: clusterProxyProviderFacade,
	}
}

// BatchLoadImportTenantInfo loads a list of import tasks' tenant information in a single API call.
// It depends on loading these imports' cluster information.
func (s *ImportHelperService) BatchLoadImportTenantInfo(
	ctx context.Context,
	impTasks []*entity.ImportTask,
) ([]*entity.ImportTask, error) {
	logger := utils.GetWorkerLoggerFromContext(ctx)
	if len(impTasks) == 0 {
		// avoid listing clusters with no cluster ID specified, in this case, it will iterate all the clusters
		return nil, nil
	}
	var err error
	if impTasks, err = batchLoadImportClusterInfo(ctx, impTasks); err != nil {
		logger.Error("batch load cluster info error", zap.Error(err))
	}
	if len(impTasks) == 0 {
		return nil, nil
	}
	tenantIDs := make([]uint64, len(impTasks))
	for i, imp := range impTasks {
		tenantIDs[i] = imp.Cluster.OrganizationId
	}
	tenantMap, err := s.accountFacade.ListOrgMapByIDs(ctx, tenantIDs)
	if err != nil {
		return nil, errors.Annotate(err, "list tenant info map by ids failed")
	}
	var missingTenantImportIDs []string
	var imps []*entity.ImportTask
	for _, imp := range impTasks {
		org, ok := tenantMap[imp.Cluster.OrganizationId]
		if !ok {
			missingTenantImportIDs = append(missingTenantImportIDs, imp.ID)
			continue
		}
		imp.Org = org
		imps = append(imps, imp)
	}
	if len(missingTenantImportIDs) > 0 {
		return imps, errors.Errorf("cannot find tenant information for some import IDs: %v", missingTenantImportIDs)
	}
	return imps, nil
}

// GetImportsToSync pass nil for all tenant or with a specified tenant.
func (s *ImportHelperService) GetImportsToSync(ctx context.Context, tenantID *entity.ID) ([]*entity.ImportTask, error) {
	logger := utils.GetWorkerLoggerFromContext(ctx)
	var imports []*entity.Import
	imports, bizErr := s.importRepo.ListImportsByFilter(ctx, &entity.ImportFilter{
		CleanUp: aws.Bool(false),
	})
	if bizErr != nil {
		return nil, errors.Annotate(bizErr, "failed to fetch imports")
	}
	results := make([]*entity.ImportTask, len(imports))
	for i := range imports {
		results[i] = &entity.ImportTask{
			Import: imports[i],
		}
	}
	var err error
	if results, err = batchLoadImportClusterInfo(ctx, results); err != nil {
		logger.Error("batch load cluster info error", zap.Error(err))
	}
	if len(results) == 0 {
		return nil, nil
	}

	if tenantID == nil {
		return results, nil
	}

	var filteredResults []*entity.ImportTask
	for _, imp := range results {
		if imp.Cluster.OrganizationId == uint64(*tenantID) {
			filteredResults = append(filteredResults, imp)
		}
	}
	return filteredResults, nil
}

// GetFailImportsInRecent pass nil for all tenant or with a specified tenant.
func (s *ImportHelperService) GetFailImportsInRecent(
	ctx context.Context,
	tenantID *entity.ID,
	recentDur time.Duration,
) ([]*entity.ImportTask, error) {
	logger := utils.GetWorkerLoggerFromContext(ctx)
	imports, bizErr := s.importRepo.ListImportsByFilter(ctx, &entity.ImportFilter{
		CompletedAtAfter: aws.Time(time.Now().Add(-recentDur)),
		Statuses:         []constant.ImportStatus{constant.ImportStatusFailed},
	})
	if bizErr != nil {
		return nil, errors.Annotate(bizErr, "failed to fetch imports")
	}
	results := make([]*entity.ImportTask, len(imports))
	for i := range imports {
		results[i] = &entity.ImportTask{
			Import: imports[i],
		}
	}
	var err error
	if results, err = s.BatchLoadImportTenantInfo(ctx, results); err != nil {
		logger.Error("batch load tenant info error", zap.Error(err))
	}
	if len(results) == 0 {
		return nil, nil
	}
	if tenantID == nil {
		return results, nil
	}
	var filteredResults []*entity.ImportTask
	for _, imp := range results {
		if imp.Cluster.OrganizationId == uint64(*tenantID) {
			filteredResults = append(filteredResults, imp)
		}
	}
	return filteredResults, nil
}

type WarningResp struct {
	Tables [][]byte `json:"table_name"`
}

// CheckWarningsByTables checks warnings for multiple tables in one query.
func (s *ImportHelperService) CheckWarningsByTables(
	ctx context.Context,
	imp *entity.ImportTask,
	tableNames []string,
) (map[string]string, error) {
	result := make(map[string]string, len(tableNames))
	if len(tableNames) == 0 {
		return result, nil
	}
	clusterID, err := strconv.ParseUint(imp.ClusterID, 10, 64)
	if err != nil {
		return result, err
	}
	passwd, bizErr := s.accountFacade.GetCloudAdminPasswd(ctx, clusterID)
	if bizErr != nil {
		return result, bizErr
	}
	password, err := s.encryptorFacade.DecryptBase64(passwd)
	if err != nil {
		return result, err
	}
	operator := &clusterproxy.ClusterOperator{
		Username:   constant.CloudAdminUsername,
		AuthMethod: clusterproxy.AuthMethodPassword,
		Credential: password,
	}
	proxy, err := s.clusterProxyProviderFacade.GetClusterProxy(ctx, imp.Cluster, operator)
	if err != nil {
		return result, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to get cluster proxy")
	}
	// Build IN clause for table names
	var quoted []string
	for _, t := range tableNames {
		quoted = append(quoted, fmt.Sprintf("'%s'", t))
	}
	query := fmt.Sprintf(
		"SELECT DISTINCT table_name FROM `lightning_task_info`.`conflict_error_v1` WHERE task_id=%d AND table_name IN (%s);",
		imp.TaskID,
		strings.Join(quoted, ","),
	)
	bytes, err := proxy.QuerySQL(ctx, query)
	if err != nil {
		if strings.Contains(err.Error(), "doesn't exist") {
			return result, nil
		}
		return result, errors.Annotatef(err, "failed to execute sql on cluster")
	}
	// bytes is a json like {"table_name":["t1","t2"]}
	var resp WarningResp
	err = json.Unmarshal(bytes, &resp)
	if err != nil {
		return result, errors.Annotatef(err, "failed to parse sql result")
	}
	log.WithContextL(ctx).
		Debug("query result", zap.Any("resp", resp))

	warnMsgTpl := "Table %s exists some duplicate unique keys. The conflicting records have been recorded in table `lightning_task_info`.`conflict_error_v1` for your reference."
	warnTables := map[string]struct{}{}
	for _, t := range resp.Tables {
		warnTables[string(t)] = struct{}{}
	}
	for _, t := range tableNames {
		if _, ok := warnTables[t]; ok {
			result[t] = fmt.Sprintf(warnMsgTpl, t)
		} else {
			result[t] = ""
		}
	}
	return result, nil
}

func (s *ImportHelperService) CreateTargetTable(
	ctx context.Context,
	cluster *tidbmgmtservicev1.Cluster,
	localImportCfg *entity.LocalImportCfg,
	columns []*entity.ColumnInfo,
	pkColumns []string,
) bizerr.BizError {
	if len(columns) == 0 {
		return nil
	}
	logger := utils.GetWorkerLoggerFromContext(ctx)
	if utf8.RuneCountInString(localImportCfg.TargetDB) > mysql.MaxDatabaseNameLength {
		return bizerr.New(
			bizerr.ErrCodeParams,
			"The schema name(%s)'s length should be less than or equal to %d",
			localImportCfg.TargetDB,
			mysql.MaxDatabaseNameLength,
		)
	}
	if utf8.RuneCountInString(localImportCfg.TargetTable) > mysql.MaxTableNameLength {
		return bizerr.New(
			bizerr.ErrCodeParams,
			"The table name(%s)'s length should be less than or equal to %d",
			localImportCfg.TargetTable,
			mysql.MaxTableNameLength,
		)
	}
	if len(pkColumns) > mysql.MaxKeyParts {
		return bizerr.New(
			bizerr.ErrCodeParams,
			"The number of primary key columns(%d) should be less than or equal to %d",
			len(pkColumns),
			mysql.MaxKeyParts,
		)
	}
	cols := make(map[string]struct{})
	for _, col := range columns {
		if utf8.RuneCountInString(col.ColumnName) > mysql.MaxColumnNameLength {
			return bizerr.New(
				bizerr.ErrCodeParams,
				"The column name(%s)'s length should be less than or equal to %d",
				col.ColumnName,
				mysql.MaxColumnNameLength,
			)
		}
		if _, ok := cols[col.ColumnName]; ok {
			return bizerr.New(bizerr.ErrCodeParams, "There are two same column names(%s).", col.ColumnName)
		}
		cols[col.ColumnName] = struct{}{}
	}

	createDBSQL := generateCreateDatabaseSQL(localImportCfg.TargetDB)
	createTableSQL, bizErr := generateCreateTableSQL(
		logger,
		localImportCfg.TargetDB,
		localImportCfg.TargetTable,
		columns,
		pkColumns,
	)
	if bizErr != nil {
		logger.Error("generate create table SQL error", zap.Error(bizErr), zap.Reflect("columns", columns))
		return bizErr
	}

	passwd, bizErr := s.accountFacade.GetCloudAdminPasswd(ctx, cluster.ClusterId)
	if bizErr != nil {
		logger.Error("get cloud admin password error", zap.Error(bizErr))
		return bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}
	password, err := s.encryptorFacade.DecryptBase64(passwd)
	if err != nil {
		logger.Error("decrypt password error", zap.Error(err))
		return bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}

	operator := &clusterproxy.ClusterOperator{
		Username:   constant.CloudAdminUsername,
		AuthMethod: clusterproxy.AuthMethodPassword,
		Credential: password,
	}

	proxy, err := s.clusterProxyProviderFacade.GetClusterProxy(ctx, cluster, operator)
	if err != nil {
		logger.Error("get cluster proxy error", zap.Error(err))
		return bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to get cluster proxy")
	}

	if err := proxy.ExecuteSQL(ctx, createDBSQL); err != nil {
		if strings.Contains(
			strings.ToLower(err.Error()),
			strings.ToLower(errno.MySQLErrName[errno.ErrResourceGroupThrottled].Raw),
		) {
			return bizerr.New(bizerr.ErrCodePermission, ErrExceededResourceGroup)
		}
		logger.Error("execute SQL to create database error", zap.Error(err), zap.String("createDBSQL", createDBSQL))
		return errorcode.MarkClusterSQLErrorUserActionable(err)
	}

	if err := proxy.ExecuteSQL(ctx, createTableSQL); err != nil {
		if strings.Contains(err.Error(), "Specified key was too long") {
			return bizerr.New(
				bizerr.ErrCodeParams,
				"Specified key was too long, max key length is 3072 bytes. Please specify a valid primary key.",
			)
		}
		if strings.Contains(
			strings.ToLower(err.Error()),
			strings.ToLower(errno.MySQLErrName[errno.ErrResourceGroupThrottled].Raw),
		) {
			return bizerr.New(bizerr.ErrCodePermission, ErrExceededResourceGroup)
		}
		logger.Error("execute SQL to create table error", zap.Error(err), zap.String("createTableSQL", createTableSQL))
		return errorcode.MarkClusterSQLErrorUserActionable(err)
	}
	return nil
}

func (s *ImportHelperService) GetTableColumns(
	ctx context.Context,
	cluster *tidbmgmtservicev1.Cluster,
	schema string,
	table string,
) (*agent.Table, bizerr.BizError) {
	logger := log.WithContextL(ctx)
	passwd, bizErr := s.accountFacade.GetCloudAdminPasswd(ctx, cluster.ClusterId)
	if bizErr != nil {
		logger.Error("get cloud admin password error", zap.Error(bizErr))
		return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}
	password, err := s.encryptorFacade.DecryptBase64(passwd)
	if err != nil {
		logger.Error("decrypt password error", zap.Error(err))
		return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}

	operator := &clusterproxy.ClusterOperator{
		Username:   constant.CloudAdminUsername,
		AuthMethod: clusterproxy.AuthMethodPassword,
		Credential: password,
	}
	proxy, err := s.clusterProxyProviderFacade.GetClusterProxy(ctx, cluster, operator)
	if err != nil {
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to get cluster proxy")
	}

	sql := fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`;", utils.EscapeBackticks(schema), utils.EscapeBackticks(table))
	res, bizErr := proxy.QuerySQL(ctx, sql)
	if bizErr != nil {
		var clusterProxyErr *facade.ClusterProxyErr
		if errors.As(bizErr, &clusterProxyErr) {
			if clusterProxyErr.SQLErrNo == mysql.ErrNoSuchTable {
				return nil, bizerr.New(errorcode.ErrCodeTableNotExist, "table %s.%s does not exist", schema, table)
			}
		}
		logger.Error("execute the SQL on target DB by agent has errors.", zap.Error(bizErr), zap.String("sql", sql))
		return nil, bizErr
	}

	sqlResult := showCreateTableResult{}
	err = json.Unmarshal(res, &sqlResult)
	if err != nil {
		logger.Error(
			"failed to unmarshal the result of show create table",
			zap.Error(err),
			zap.String("resString", string(res)),
		)
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to unmarshal the result of show create table")
	}

	var createTableSql string
	for _, bytes := range sqlResult.CreateTable {
		createTableSql += string(bytes)
	}
	stmt, err := parser.New().ParseOneStmt(createTableSql, "", "")
	if err != nil {
		logger.Error("parse create table sql error", zap.Error(err), zap.String("createTableSql", createTableSql))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to parse create table SQL")
	}
	createTableStmt, ok := stmt.(*ast.CreateTableStmt)
	if !ok {
		logger.Error("the statement is not create table", zap.String("createTableSql", createTableSql))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "the statement is not create table: %s", createTableSql)
	}
	result := &agent.Table{
		SchemaName: schema,
		TableName:  table,
	}
	for _, col := range createTableStmt.Cols {
		result.RowHeaders = append(result.RowHeaders, col.Name.Name.O)
		result.RowDataTypes = append(result.RowDataTypes, strings.ToUpper(col.Tp.CompactStr()))
		result.ColumnIsPk = append(result.ColumnIsPk, columnIsPkuk(col, createTableStmt.Constraints))
		result.ColumnAutoGen = append(result.ColumnAutoGen, columnIsAutoGen(col))
	}
	return result, nil
}