package main

import (
	"context"
	"net"
	"sync"
	"time"

	"go.uber.org/zap"
)

// readinessProbe verifies that a backend dependency the proxy needs to serve
// traffic is reachable. It returns nil when the dependency is ready.
type readinessProbe func(ctx context.Context) error

// tcpReachableProbe returns a readinessProbe that verifies a TCP connection can
// be established to addr (host:port) within timeout. It is the minimal,
// credential-free check that the MySQL backend (e.g. TiProxy) is reachable
// before the proxy starts accepting CQL traffic.
func tcpReachableProbe(addr string, timeout time.Duration) readinessProbe {
	return func(ctx context.Context) error {
		dialer := net.Dialer{Timeout: timeout}
		conn, err := dialer.DialContext(ctx, "tcp", addr)
		if err != nil {
			return err
		}
		return conn.Close()
	}
}

// warmUp runs probe immediately and then every interval until it succeeds,
// returning nil on the first success. It returns ctx.Err() if ctx is cancelled
// before the probe succeeds, so a never-ready backend cannot block forever.
func warmUp(ctx context.Context, probe readinessProbe, interval time.Duration, log *zap.Logger) error {
	if log == nil {
		log = zap.NewNop()
	}
	for attempt := 1; ; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}
		if err := probe(ctx); err == nil {
			return nil
		} else if ctxErr := ctx.Err(); ctxErr != nil {
			// Cancelled during the probe (e.g. shutdown): abort without logging
			// a misleading "retrying".
			return ctxErr
		} else {
			log.Warn("warmup probe failed, retrying",
				zap.Int("attempt", attempt),
				zap.Duration("interval", interval),
				zap.Error(err))
		}
		timer := time.NewTimer(interval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}
}

// readiness tracks whether the proxy should advertise itself as ready on /ready.
//
// Ready == warmed && !draining. The proxy is only ready once warmup has verified
// its backend dependencies (warmed) and a drain has not started (draining).
// Draining is terminal: once a drain begins the proxy never reports ready again,
// even if a late warmup completes afterwards. State transitions are serialized
// under mu so a warmup that races a drain cannot mark a draining pod ready.
type readiness struct {
	mu       sync.Mutex
	warmed   bool
	draining bool
}

// Ready reports whether /ready should return 200.
func (r *readiness) Ready() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.warmed && !r.draining
}

// MarkWarm records that warmup completed successfully. It is a no-op once the
// proxy is draining (drain is terminal). It returns whether the proxy is now
// ready, so the caller can avoid logging a misleading "ready" after a drain.
func (r *readiness) MarkWarm() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.draining {
		return false
	}
	r.warmed = true
	return true
}

// Draining reports whether the proxy has entered drain mode.
func (r *readiness) Draining() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.draining
}

// MarkDraining records that the proxy has entered drain mode.
func (r *readiness) MarkDraining() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.draining = true
}