package cluster

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

	_ "github.com/go-sql-driver/mysql"
	"github.com/onsi/gomega"
	"github.com/pingcap/errors"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
)

func endpointConnected(ctx context.Context, cluster *infra.Cluster) error {
	dsn := DSN(cluster.Status.KeyspaceToken, *cluster.Status.SqlEndpoints.Public)
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return errors.Annotatef(err, "failed to open db %q", dsn)
	}
	defer db.Close() //nolint:errcheck // ignore error

	row := db.QueryRowContext(ctx, "select 1")
	if row.Err() != nil {
		return errors.Annotatef(row.Err(), "failed to run query")
	}

	return nil
}

func waitTiDBReady(ctx context.Context, cluster *infra.Cluster, timeout, interval time.Duration) {
	gomega.Eventually(func() error {
		return endpointConnected(ctx, cluster)
	}, timeout, interval).Should(gomega.Succeed(), "SQL endpoint should be ready")
}

// DSN returns a DSN string for a TiDB cluster.
func DSN(keyspaceToken string, endpoint infra.PublicStatus) string {
	return fmt.Sprintf("%s.root@tcp(%s:%d)/?tls=preferred", keyspaceToken, endpoint.Host, endpoint.Port)
}