package executor

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

	"go.uber.org/zap"
)

const tidbPartitionAttributeValue = `region_policy=exclusive`

func isTiDBVersionString(version string) bool {
	return strings.Contains(strings.ToUpper(strings.TrimSpace(version)), "TIDB")
}

func buildPartitionAttributeSQL(schema, table, partition, attributes string) string {
	return fmt.Sprintf(
		"ALTER TABLE %s.%s PARTITION %s ATTRIBUTES=%q",
		quoteMySQLIdentifier(schema),
		quoteMySQLIdentifier(table),
		quoteMySQLIdentifier(partition),
		attributes,
	)
}

func (e *Executor) cachedBackendIsTiDB() (bool, bool) {
	if e == nil {
		return false, false
	}
	e.backendInfoMu.RLock()
	defer e.backendInfoMu.RUnlock()
	return e.backendIsTiDB, e.backendIsTiDBKnown
}

func (e *Executor) setCachedBackendIsTiDB(isTiDB bool) {
	if e == nil {
		return
	}
	e.backendInfoMu.Lock()
	defer e.backendInfoMu.Unlock()
	e.backendIsTiDB = isTiDB
	e.backendIsTiDBKnown = true
}

func detectTiDBBackend(ctx context.Context, conn *sql.Conn) (bool, error) {
	if conn == nil {
		return false, fmt.Errorf("connection is required")
	}
	var version sql.NullString
	if err := conn.QueryRowContext(ctx, "SELECT VERSION()").Scan(&version); err != nil {
		return false, err
	}
	return isTiDBVersionString(version.String), nil
}

func (e *Executor) ensureBackendIsTiDB(ctx context.Context, user, password string, conn *sql.Conn) (bool, error) {
	if cached, ok := e.cachedBackendIsTiDB(); ok {
		return cached, nil
	}

	var (
		release func()
		err     error
	)
	if conn == nil {
		conn, release, err = e.getConn(ctx, user, password)
		if err != nil {
			return false, err
		}
		if release != nil {
			defer release()
		}
	}

	isTiDB, err := detectTiDBBackend(ctx, conn)
	if err != nil {
		return false, err
	}
	e.setCachedBackendIsTiDB(isTiDB)
	return isTiDB, nil
}

func loadPartitionNames(ctx context.Context, conn *sql.Conn, schema, table string) ([]string, error) {
	if conn == nil {
		return nil, fmt.Errorf("connection is required")
	}
	rows, err := conn.QueryContext(ctx, `
		SELECT PARTITION_NAME
		FROM INFORMATION_SCHEMA.PARTITIONS
		WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND PARTITION_NAME IS NOT NULL
		ORDER BY PARTITION_ORDINAL_POSITION
	`, schema, table)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	partitions := make([]string, 0)
	for rows.Next() {
		var name sql.NullString
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		if strings.TrimSpace(name.String) == "" {
			continue
		}
		partitions = append(partitions, name.String)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return partitions, nil
}

func applyPartitionAttributes(ctx context.Context, conn *sql.Conn, schema, table string, partitions []string, attributes string) error {
	if conn == nil {
		return fmt.Errorf("connection is required")
	}
	for _, partition := range partitions {
		if strings.TrimSpace(partition) == "" {
			continue
		}
		sqlText := buildPartitionAttributeSQL(schema, table, partition, attributes)
		if _, err := conn.ExecContext(ctx, sqlText); err != nil {
			return err
		}
	}
	return nil
}

func (e *Executor) applyTiDBPartitionAttributes(ctx context.Context, user, password, schema, table string) error {
	if e == nil || schema == "" || table == "" {
		return nil
	}

	conn, release, err := e.getConn(ctx, user, password)
	if err != nil {
		return err
	}
	if release != nil {
		defer release()
	}

	isTiDB, err := e.ensureBackendIsTiDB(ctx, user, password, conn)
	if err != nil {
		return err
	}
	if !isTiDB {
		return nil
	}

	partitions, err := loadPartitionNames(ctx, conn, schema, table)
	if err != nil {
		return err
	}
	if len(partitions) == 0 {
		return nil
	}
	return applyPartitionAttributes(ctx, conn, schema, table, partitions, tidbPartitionAttributeValue)
}

func (e *ddlExecutor) launchPostCreateTableTasks(user, password, schema, table string) {
	if e == nil || e.Executor == nil || schema == "" || table == "" {
		return
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()

		if err := e.applyTiDBPartitionAttributes(ctx, user, password, schema, table); err != nil && e.logger != nil {
			e.logger.Warn("post-create-table async task failed",
				zap.String("keyspace", schema),
				zap.String("table", table),
				zap.Error(err))
		}
	}()
}