package probe

import (
	"context"
	"database/sql"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/go-logr/logr"
	_ "github.com/go-sql-driver/mysql"
	"github.com/pkg/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/router"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	ctrl "sigs.k8s.io/controller-runtime"
)

const (
	ProbeTiDBRetry       = 2 * time.Second
	ProbeTiDBTimeout     = 6 * time.Second
	ProbeTiDBInterval    = 20 * time.Second
	ProbeLatencyInterval = 15 * time.Second
	// ProbeExitGracePeriod: before exiting due to ConnCount<=0, re-check after this delay to avoid
	// race with short-lived connections (new conn may have just called Inc while we saw 0).
	ProbeExitGracePeriod = 3 * time.Second
)

func NewTiDBProber(
	ctx context.Context,
	rb *router.RouterBundle,
	regionalInfra infra.RegionalInfra,
	infraApi infra.InfraAPI,
	tidbController tidb.Controller,
) *TiDBProber {
	return &TiDBProber{
		ctx:            ctx,
		routerBundle:   rb,
		regionalInfra:  regionalInfra,
		infraApi:       infraApi,
		tidbController: tidbController,
		cc:             make(map[string]bool),

		logger: ctrl.Log.WithName("probeTiDBServer"),
	}
}

type TiDBProber struct {
	ctx            context.Context
	routerBundle   *router.RouterBundle
	regionalInfra  infra.RegionalInfra
	infraApi       infra.InfraAPI
	tidbController tidb.Controller

	wg     sync.WaitGroup
	logger logr.Logger

	mu sync.Mutex
	cc map[string]bool // token -> whether probe goroutine is running for this token
}

func (s *TiDBProber) Stop() {
	s.wg.Wait()
	s.logger.Info("probe TiDB server stopped")
}

// Inc ensures the probe goroutine is running for this token. Called from OnHandshake when backend
// handshake succeeds. Only one probe goroutine runs per token; multiple connections just keep it alive.
func (s *TiDBProber) Inc(token, clusterID, connID string) int64 {
	s.mu.Lock()
	if s.cc[token] {
		s.mu.Unlock()
		return 0
	}
	s.cc[token] = true
	s.mu.Unlock()
	go s.startProbeTiDBPod(token, clusterID)
	return 0
}

//nolint:lll // ignore
func (s *TiDBProber) startProbeTiDBPod(token, clusterID string) {
	s.wg.Add(1)
	defer s.wg.Done()
	defer func() {
		s.mu.Lock()
		delete(s.cc, token)
		s.mu.Unlock()
	}()

	logger := s.logger.WithValues("token", token, "cluster_id", clusterID)
	logger.Info("probe TiDB server start")

	t := time.NewTicker(ProbeTiDBInterval)
	defer t.Stop()

	// Require two consecutive ticks with ConnCount<=0 before exiting, so short-lived connections
	// (connect/disconnect between ticks) don't cause probe to exit and immediately restart on next Inc.
	lastTickWasZero := false

	for {
		select {
		case <-s.ctx.Done():
			return
		case <-t.C:
			if s.routerBundle.GetConnCountByToken(token) <= 0 {
				// Re-check after grace period to avoid race: new conn may have just called Inc.
				select {
				case <-s.ctx.Done():
					return
				case <-time.After(ProbeExitGracePeriod):
				}
				if s.routerBundle.GetConnCountByToken(token) <= 0 {
					if lastTickWasZero {
						logger.Info("probe TiDB server stop, no active connection")
						return
					}
					lastTickWasZero = true
					continue
				}
				lastTickWasZero = false
			} else {
				lastTickWasZero = false
			}
			clusterStatus, err := s.getAwsFreeTierClusterStatus(token)
			if err != nil {
				logger.Error(err, "probe TiDB server stop, get cluster status failed")
				return
			}
			logger.V(1).Info("probe job start")
			start := time.Now()
			for _, ip := range s.routerBundle.GetBackendIPs(token) {
				var err error
				status := config.CONNECTION_STATUS_CANCEL
				if time.Since(start) < ProbeTiDBInterval {
					status, err = s.doProbe(token, ip)
				}
				if err != nil || status != config.CONNECTION_STATUS_NORMAL {
					logger.Info(fmt.Sprintf("probe job encounter error, detail: target=%s, status=%s, duration=%f, err=%v", ip, status, time.Since(start).Seconds(), err))
				}
				metrics.ConnectionSLACounter.WithLabelValues(clusterID, token, clusterStatus.TidbPool, clusterStatus.RuLimitMode, "", "probe", status, clusterStatus.ServicePlan).Inc()
				metrics.ProbeTiDBHistogram.WithLabelValues(clusterID, token, clusterStatus.TidbPool, clusterStatus.RuLimitMode, status).Observe(time.Since(start).Seconds())
			}
			logger.V(1).Info("probe job done")
		}
	}
}

func (s *TiDBProber) doProbe(token, ip string) (status string, err error) {
	logger := s.logger.WithValues("token", token).WithValues("ip", ip)
	defer func() {
		if err != nil {
			logger.Error(err, "doProbe: probe tidb pod failed")
			status = config.CONNECTION_STATUS_INNER_ERROR
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), ProbeTiDBTimeout)
	defer cancel()
	resp, err := s.HealthCheck(ctx, ip)
	if err != nil {
		return "", err
	}
	if resp.Status == tidb.HealthStatusStandby || resp.Token != token {
		logger.Info("[Warn] doProbe: target server mismatch",
			"requestToken", token,
			"responseToken", resp.Token,
		)
		return config.CONNECTION_STATUS_MISSMATCH, nil
	}
	if resp.Status == tidb.HealthStatusUp {
		return config.CONNECTION_STATUS_NORMAL, nil
	}
	logger.Error(errors.New(resp.Error), "doProbe: tidb server down")
	return config.CONNECTION_STATUS_TIDB_ERROR, nil
}

func (s *TiDBProber) HealthCheck(ctx context.Context, ip string) (*tidb.HealthCheckResp, error) {
	host := fmt.Sprintf("%s:10080", ip)
	resp, err := s.tidbController.HealthCheck(ctx, host)
	if err != nil && !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
		var opError *net.OpError
		// if any network error, retry again(tidb may be in restarting)
		if errors.As(err, &opError) {
			time.Sleep(ProbeTiDBRetry)
			resp, err = s.tidbController.HealthCheck(ctx, host)
		}
	}
	return resp, err
}

/* monitor latency with a targeted cluster */
var (
	ConcurrencyNum = 10
	RecordsCount   = 30
	STEP           = RecordsCount / ConcurrencyNum
)

func prepareDBData(db *sql.DB) error {
	//nolint: lll // ignore long line length due to kubebuilder validation rules
	_, err := db.Exec("CREATE TABLE IF NOT EXISTS serverless_probe (id INT(11) NOT NULL AUTO_INCREMENT PRIMARY KEY, updated_at TIMESTAMP)")
	if err != nil {
		return err
	}
	stmt, err := db.Prepare("INSERT INTO serverless_probe (updated_at) VALUES (?)")
	if err != nil {
		return err
	}
	defer stmt.Close() //nolint:errcheck // ignore error
	for i := 0; i < RecordsCount; i++ {
		if _, err = stmt.Exec(time.Now().UTC()); err != nil {
			return err
		}
	}
	return nil
}

func prepareTTLTable(db *sql.DB) error {
	_, err := db.Exec("CREATE TABLE IF NOT EXISTS serverless_probe_ttl_1h (" +
		"id INT(11) NOT NULL AUTO_INCREMENT PRIMARY KEY," +
		"message varchar(256) NOT NULL," +
		"created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP," +
		"INDEX idx_created_at (created_at)" +
		") TTL = `created_at` + INTERVAL 1 HOUR;")
	if err != nil {
		return err
	}
	stmt, err := db.Prepare("INSERT INTO serverless_probe_ttl_1h (message) VALUES (?)")
	if err != nil {
		return err
	}
	_, err = stmt.Exec("test_ttl")
	stmt.Close() //nolint:errcheck // ignore error
	return err
}

func initDB(logger logr.Logger, dsn string) (db *sql.DB, err error) {
	// dsn = "<user>:<password>@tcp(host:4000)/test?tls=tidb"
	// host = "gateway01.us-east-1.staging.shared.aws.tidbcloud.com"
	dsn = strings.Split(dsn, "tls=tidb")[0]
	startTime := time.Now()
	host := strings.Split(strings.Split(dsn, "tcp(")[1], ":4000")[0]
	user := strings.Split(dsn, ":")[0]
	logger.Info("init db start..", "origin_host", host, "user", user, "dsn", dsn)
	defer func() {
		if err != nil {
			metrics.ProbeLatencyHistogram.WithLabelValues(user, "INIT_DB", "FAIL").Observe(time.Since(startTime).Seconds())
			logger.Error(err, "init db failed.")
		}
		logger.Info("init db success....")
	}()
	/*
		// public gateway force tls
		dsn = dsn + "&tls=tidb"
		mysql.RegisterTLSConfig("tidb", &tls.Config{ //nolint:errcheck // ignore error
			MinVersion: tls.VersionTLS12,
			ServerName: host,
		})
	*/
	/* change host to inner gateway svc */
	dsn = strings.Replace(dsn, host, "tidb-gateway-inner.tidb-admin", 1)
	db, err = sql.Open("mysql", dsn+"&parseTime=true")
	if err != nil {
		return nil, err
	}
	db.SetMaxIdleConns(ConcurrencyNum)
	db.SetMaxOpenConns(ConcurrencyNum)
	db.SetConnMaxLifetime(time.Minute * 10) //nolint:mnd // ignore
	if err = db.Ping(); err != nil {
		db.Close() //nolint:errcheck // ignore error
		return nil, err
	}
	if err = prepareDBData(db); err != nil {
		db.Close() //nolint:errcheck // ignore error
		return nil, err
	}
	if err = prepareTTLTable(db); err != nil {
		db.Close() //nolint:errcheck // ignore error
		return nil, err
	}
	return db, nil
}

func RunProbeLatency(ctx context.Context) {
	ticker := time.NewTicker(ProbeLatencyInterval)
	defer ticker.Stop()
	var (
		err    error
		logger = ctrl.Log.WithName("probeTiDBLatency")
		dsnM   = make(map[string]string)
		dbs    = make(map[string]*sql.DB)
	)

	for _, dsn := range strings.Split(os.Getenv("PROBE_SERVERLESS_DSN_SECRET"), "#") {
		if len(dsn) > 0 {
			user := strings.Split(dsn, ":")[0]
			dsnM[user] = dsn
			dbs[user] = nil
		}
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			for user, dsn := range dsnM {
				if dbs[user] == nil {
					if dbs[user], err = initDB(logger, dsn); err != nil {
						logger.Error(err, "db is nil, reinit failed, skip probe. dsn = "+dsn)
						continue
					}
				}
				probeLatency(ctx, logger, dbs[user], user)
			}
		}
	}
}

func logTraceForSlowProbe(logger logr.Logger, index int, rows *sql.Rows, elapsed time.Duration) {
	if rows != nil {
		defer rows.Close()             //nolint:errcheck // ignore error
		if elapsed.Seconds() >= 0.05 { //nolint:mnd // hardcoded 50ms to now
			var operation string
			var startTs string
			var duration string
			for rows.Next() {
				if err := rows.Scan(&operation, &startTs, &duration); err != nil {
					logger.Error(err, "Read trace failed.")
					return
				}
				logger.Info("Slow Probe", "index", index, "operation", operation, "startTs", startTs, "duration", duration)
			}
		}
	}
}

func probeLatency(ctx context.Context, log logr.Logger, db *sql.DB, user string) {
	logger := log.WithValues("user", user)
	startTime := time.Now()
	logger.Info("probe job start")
	defer func() {
		logger.Info("probe job end", "elapse", time.Since(startTime).Seconds())
	}()

	var wg sync.WaitGroup
	for i := 1; i <= ConcurrencyNum; i++ {
		wg.Add(1)
		go func(offset int) {
			defer wg.Done()
			s := time.Now()
			conn, err := db.Conn(ctx)
			if err != nil {
				metrics.ProbeLatencyHistogram.WithLabelValues(user, "GET_CONN_POOL", "FAIL").Observe(time.Since(s).Seconds())
				logger.Error(err, "probe job without conn in pool.")
				return
			}
			metrics.ProbeLatencyHistogram.WithLabelValues(user, "GET_CONN_POOL", "OK").Observe(time.Since(s).Seconds())
			defer conn.Close() //nolint:errcheck // ignore error

			var rows *sql.Rows
			var elapsed time.Duration
			for j := (offset-1)*STEP + 1; j <= offset*STEP; j++ {
				status := "OK"
				s = time.Now()
				_, err = conn.ExecContext(ctx, "UPDATE serverless_probe SET updated_at = ? WHERE id = ?", startTime.UTC(), j) //nolint:lll // ignore error
				if err != nil {
					status = "FAIL"
					logger.Error(err, "probe job update db failed.")
				}
				elapsed = time.Since(s)
				metrics.ProbeLatencyHistogram.WithLabelValues(user, "UPDATE", status).Observe(elapsed.Seconds())

				status = "OK"
				s = time.Now()
				rows, err = conn.QueryContext(ctx, "TRACE SELECT updated_at FROM serverless_probe WHERE id = ?", j)
				if err != nil {
					status = "FAIL"
					logger.Error(err, "probe job query db failed.")
				}
				elapsed = time.Since(s)
				metrics.ProbeLatencyHistogram.WithLabelValues(user, "SELECT", status).Observe(elapsed.Seconds())
				logTraceForSlowProbe(logger, offset, rows, elapsed)

				/* insert ttl table with autoincrementing id */
				status = "OK"
				s = time.Now()
				_, err = conn.ExecContext(ctx, "INSERT INTO serverless_probe_ttl_1h (message) VALUES (?)", startTime.UTC().String()) //nolint:lll // ignore error
				if err != nil {
					status = "FAIL"
					logger.Error(err, "probe job update ttl table failed.")
				}
				elapsed = time.Since(s)
				metrics.ProbeLatencyHistogram.WithLabelValues(user, "TTL_INSERT", status).Observe(elapsed.Seconds())

				/* ttl expire check */
				status = "OK"
				s = time.Now()
				var expiredDataCount int
				err = conn.QueryRowContext(ctx, "select count(*) from serverless_probe_ttl_1h where created_at < ?", startTime.Add(-6*time.Hour).UTC()).Scan(&expiredDataCount) //nolint:lll // ignore error
				if err != nil {
					status = "FAIL"
					logger.Error(err, "probe job ttl table expired data check failed.")
				}
				if expiredDataCount != 0 {
					status = "FAIL"
					logger.Error(fmt.Errorf("expired data count = %d", expiredDataCount),
						"probe job ttl table expired data check failed.")
				}
				elapsed = time.Since(s)
				metrics.ProbeLatencyHistogram.WithLabelValues(user, "TTL_EXPIRE", status).Observe(elapsed.Seconds())
			}
		}(i)
	}
	wg.Wait()
	metrics.ProbeLatencyHistogram.WithLabelValues(user, "total_job", "-").Observe(time.Since(startTime).Seconds())
}

type basicStatus struct {
	KeyspaceToken string
	TidbPool      string
	RuLimitMode   string
	ServicePlan   string
}

func (s *TiDBProber) getAwsFreeTierClusterStatus(token string) (basicStatus basicStatus, err error) {
	infraCluster, err := s.regionalInfra.GetClusterWithToken(s.ctx, token)
	if err == nil {
		basicStatus.KeyspaceToken = infraCluster.Status.KeyspaceToken
		basicStatus.TidbPool = infraCluster.Status.TidbPool
		basicStatus.RuLimitMode = infraCluster.Status.RuLimitMode
		basicStatus.ServicePlan = infraCluster.Annotations[types.AnnoTidbServicePlan]
		return
	}
	infraBranch, err := s.infraApi.GetBranchWithToken(s.ctx, token)
	if err == nil {
		var branchStatus *cloudv1.AwsFreeTierClusterStatus
		if branchStatus, err = infraBranch.AwsFreeTierClusterStatusFromBranch(); err != nil {
			return
		}
		basicStatus.KeyspaceToken = branchStatus.KeyspaceToken
		basicStatus.TidbPool = branchStatus.TidbPool
		basicStatus.RuLimitMode = branchStatus.RuLimitMode
		basicStatus.ServicePlan = infraBranch.Annotations[types.AnnoTidbServicePlan]
		return
	}
	return
}