//go:build integration
// +build integration

package integration_test

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

func TestCreateTableAppliesTiDBPartitionAttributesAsync(t *testing.T) {
	requireProxy(t)

	session := newSession(t, sessionOptions{disableHostLookup: true})
	keyspace := uniqueKeyspace()
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	db, err := sql.Open("mysql", mysqlDSN(mysqlCfg))
	if err != nil {
		t.Fatalf("mysql connect failed: %v", err)
	}
	defer db.Close()

	if !isTiDBBackend(t, db) {
		t.Skip("backend is not TiDB")
	}

	table := "partition_attr_async"
	create := fmt.Sprintf("CREATE TABLE %s.%s (id int PRIMARY KEY, name text)", keyspace, table)
	if err := session.Query(create).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	partitions := partitionCount(t, db, keyspace, table)
	if partitions == 0 {
		t.Fatalf("expected partitioned table for %s.%s", keyspace, table)
	}

	waitFor(t, "TiDB partition attributes", 5*time.Second, func() (bool, error) {
		count, err := partitionAttributeCount(db, keyspace, table)
		if err != nil {
			return false, err
		}
		return count == partitions, nil
	})
}

func isTiDBBackend(t *testing.T, db *sql.DB) bool {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var version string
	if err := db.QueryRowContext(ctx, "SELECT VERSION()").Scan(&version); err != nil {
		t.Fatalf("version query failed: %v", err)
	}
	return strings.Contains(strings.ToUpper(version), "TIDB")
}

func partitionCount(t *testing.T, db *sql.DB, schema, table string) int {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var count int
	if err := db.QueryRowContext(ctx, `
		SELECT COUNT(*)
		FROM INFORMATION_SCHEMA.PARTITIONS
		WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND PARTITION_NAME IS NOT NULL
	`, schema, table).Scan(&count); err != nil {
		t.Fatalf("partition count query failed: %v", err)
	}
	return count
}

func partitionAttributeCount(db *sql.DB, schema, table string) (int, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var count int
	err := db.QueryRowContext(ctx, `
		SELECT COUNT(*)
		FROM INFORMATION_SCHEMA.ATTRIBUTES
		WHERE TYPE = 'key-range'
		  AND ATTRIBUTES = ?
		  AND ID LIKE CONCAT('schema/', ?, '/', ?, '/%')
	`, fmt.Sprintf("%q", "region_policy=exclusive"), schema, table).Scan(&count)
	return count, err
}