package tunserver

import (
	"context"
	"errors"
	"io"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/retry"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tunnel/rpc"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"k8s.io/apimachinery/pkg/util/wait"
)

var (
	proxyStreamDesc = grpc.StreamDesc{
		ServerStreams: true,
		ClientStreams: true,
	}

	// tunnelReadySentinelError is a sentinel error value to make stream visitor exit early.
	tunnelReadySentinelError = errors.New("")
)

type connAttempt struct {
	cancel context.CancelFunc
}

type ReadyGateway struct {
	URL          string
	Stream       grpc.ClientStream
	Conn         grpctool.PoolConn
	StreamCancel context.CancelFunc
}

func (g ReadyGateway) Done() {
	g.StreamCancel()
	g.Conn.Done()
}

// PollGatewayURLsCallback is called periodically with found kas URLs for a particular agent id.
type PollGatewayURLsCallback func(kasURLs []string)

type PollingGatewayURLQuerier interface {
	PollGatewayURLs(ctx context.Context, agentID int64, cb PollGatewayURLsCallback)
	CachedGatewayURLs(agentID int64) []string
}

type GatewayFinder interface {
	Find(ctx context.Context) (ReadyGateway, error)
}

type gatewayFinder struct {
	log                   *zap.Logger
	gatewayPool           grpctool.PoolInterface
	gatewayQuerier        PollingGatewayURLQuerier
	rpcAPI                modshared.RPCAPI
	fullMethod            string // /service/method
	ownPrivateAPIURL      string
	agentID               int64
	outgoingCtx           context.Context
	pollConfig            retry.PollConfigFactory
	foundGateway          chan ReadyGateway
	noTunnel              chan struct{}
	wg                    wait.Group
	pollCancel            context.CancelFunc
	tryNewGatewayInterval time.Duration

	mu          sync.Mutex             // protects the fields below
	connections map[string]connAttempt // gateway tunserver URL -> conn info
	gatewayURLs []string               // currently known gateway tunserver URLs for the agent id
	done        bool                   // successfully done searching
}

func NewGatewayFinder(log *zap.Logger, gatewayPool grpctool.PoolInterface, gatewayQuerier PollingGatewayURLQuerier,
	rpcAPI modshared.RPCAPI, fullMethod string, ownPrivateAPIURL string, agentID int64, outgoingCtx context.Context,
	pollConfig retry.PollConfigFactory, tryNewGatewayInterval time.Duration) GatewayFinder {
	return &gatewayFinder{
		log:                   log,
		gatewayPool:           gatewayPool,
		gatewayQuerier:        gatewayQuerier,
		rpcAPI:                rpcAPI,
		fullMethod:            fullMethod,
		ownPrivateAPIURL:      ownPrivateAPIURL,
		agentID:               agentID,
		outgoingCtx:           outgoingCtx,
		pollConfig:            pollConfig,
		tryNewGatewayInterval: tryNewGatewayInterval,
		foundGateway:          make(chan ReadyGateway),
		noTunnel:              make(chan struct{}),
		connections:           make(map[string]connAttempt),
	}
}

func (f *gatewayFinder) Find(ctx context.Context) (ReadyGateway, error) {
	defer f.wg.Wait()
	var pollCtx context.Context
	pollCtx, f.pollCancel = context.WithCancel(ctx)
	defer f.pollCancel()

	// Unconditionally connect to self ASAP.
	f.tryGatewayLocked(f.ownPrivateAPIURL) //nolint: contextcheck
	startedPolling := false
	// This flag is set when we've run out of gateway tunserver URLs to try. When a new set of URLs is received, if this is set,
	// we try to connect to one of those URLs.
	needToTryNewGateway := false

	// Timer is used to wake up the loop below after a certain amount of time has passed but there has been no activity,
	// in particular, a recently connected to gateway tunserver didn't reply with noTunnel. If it's not replying, we
	// need to try another instance if it has been discovered.
	// If, for some reason, our own private API server doesn't respond with noTunnel/startStreaming in time, we
	// want to proceed with normal flow too.
	t := time.NewTimer(f.tryNewGatewayInterval)
	defer t.Stop()
	gatewayURLsC := make(chan []string)
	f.gatewayURLs = f.gatewayQuerier.CachedGatewayURLs(f.agentID)
	done := ctx.Done()

	// Timer must have been stopped or has fired when this function is called
	tryNewGatewayWhenTimerNotRunning := func() {
		if f.tryNewGateway() { //nolint: contextcheck
			// Connected to an instance.
			needToTryNewGateway = false
			t.Reset(f.tryNewGatewayInterval)
		} else {
			// Couldn't find a gateway tunserver instance we haven't connected to already.
			needToTryNewGateway = true
			if !startedPolling {
				startedPolling = true
				// No more cached instances, start polling for gateway tunserver instances.
				f.wg.Start(func() {
					pollDone := pollCtx.Done()
					f.gatewayQuerier.PollGatewayURLs(pollCtx, f.agentID, func(gatewayURLs []string) {
						select {
						case <-pollDone:
						case gatewayURLsC <- gatewayURLs:
						}
					})
				})
			}
		}
	}

	for {
		select {
		case <-done:
			f.stopAllConnectionAttempts()
			return ReadyGateway{}, ctx.Err()
		case <-f.noTunnel:
			stopAndDrain(t)
			tryNewGatewayWhenTimerNotRunning()
		case gatewayURLs := <-gatewayURLsC:
			f.mu.Lock()
			f.gatewayURLs = gatewayURLs
			f.mu.Unlock()
			if !needToTryNewGateway {
				continue
			}
			if f.tryNewGateway() { //nolint: contextcheck
				// Connected to a new gateway instance.
				needToTryNewGateway = false
				stopAndDrain(t)
				t.Reset(f.tryNewGatewayInterval)
			}
		case <-t.C:
			tryNewGatewayWhenTimerNotRunning()
		case rt := <-f.foundGateway:
			f.stopAllConnectionAttemptsExcept(rt.URL)
			return rt, nil
		}
	}
}

func (f *gatewayFinder) tryNewGateway() bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, gatewayURL := range f.gatewayURLs {
		if _, ok := f.connections[gatewayURL]; ok {
			continue // skip gateway tunserver that we have connected to already
		}
		f.tryGatewayLocked(gatewayURL)
		return true
	}
	return false
}

func (f *gatewayFinder) tryGatewayLocked(gatewayURL string) {
	connCtx, connCancel := context.WithCancel(f.outgoingCtx)
	f.connections[gatewayURL] = connAttempt{
		cancel: connCancel,
	}
	f.wg.Start(func() {
		f.tryGatewayAsync(connCtx, connCancel, gatewayURL)
	})
}

func (f *gatewayFinder) tryGatewayAsync(ctx context.Context, cancel context.CancelFunc, gatewayURL string) {
	log := f.log.With(logz.KASURL(gatewayURL))
	noTunnelSent := false
	_ = retry.PollWithBackoff(ctx, f.pollConfig(), func(ctx context.Context) (error, retry.AttemptResult) {
		success := false

		// 1. Dial another gateway tunserver
		log.Debug("Trying tunnel")
		attemptCtx, attemptCancel := context.WithCancel(ctx)
		defer func() {
			if !success {
				attemptCancel()
				f.maybeStopTrying(gatewayURL)
			}
		}()
		gatewayConn, err := f.gatewayPool.Dial(attemptCtx, gatewayURL)
		if err != nil {
			f.rpcAPI.HandleProcessingError(log, f.agentID, "Failed to dial gateway tunserver", err)
			return nil, retry.Backoff
		}
		defer func() {
			if !success {
				gatewayConn.Done()
			}
		}()

		// 2. Open a stream to the desired service/method
		gatewayStream, err := gatewayConn.NewStream(
			attemptCtx,
			&proxyStreamDesc,
			f.fullMethod,
			grpc.ForceCodec(grpctool.RawCodecWithProtoFallback{}),
			grpc.WaitForReady(true),
		)
		if err != nil {
			f.rpcAPI.HandleProcessingError(log, f.agentID, "Failed to open a new stream to gateway tunserver", err)
			return nil, retry.Backoff
		}

		// 3. Wait for the gateway tunserver to say it's ready to start streaming i.e. has a suitable tunnel to an agent
		err = rpc.GatewayResponseVisitor().Visit(gatewayStream,
			grpctool.WithCallback(noTunnelFieldNumber, func(noTunnel *rpc.GatewayResponse_NoTunnel) error {
				trace.SpanFromContext(gatewayStream.Context()).AddEvent("No tunnel") //nolint: contextcheck
				if !noTunnelSent {                                                   // send only once
					noTunnelSent = true
					// Let Find() know there is no tunnel available from that gateway tunserver instantaneously.
					// A tunnel may still be found when a suitable agent connects later, but none available immediately.
					select {
					case <-attemptCtx.Done():
					case f.noTunnel <- struct{}{}:
					}
				}
				return nil
			}),
			grpctool.WithCallback(tunnelReadyFieldNumber, func(tunnelReady *rpc.GatewayResponse_TunnelReady) error { //nolint:contextcheck
				trace.SpanFromContext(gatewayStream.Context()).AddEvent("Ready")
				return tunnelReadySentinelError
			}),
			grpctool.WithNotExpectingToGet(codes.Internal, headerFieldNumber, messageFieldNumber, trailerFieldNumber, errorFieldNumber),
		)
		switch err { //nolint:errorlint
		case nil:
			// Gateway tunserver closed the connection cleanly, perhaps it's been open for too long
			return nil, retry.ContinueImmediately
		case tunnelReadySentinelError:
			// fallthrough
		default:
			f.rpcAPI.HandleProcessingError(log, f.agentID, "RecvMsg(GatewayResponse)", err)
			return nil, retry.Backoff
		}

		// 4. Check if another goroutine has found a suitable tunnel already
		f.mu.Lock() // Ensure only one gateway tunserver gets StartStreaming message
		if f.done {
			f.mu.Unlock()
			return nil, retry.Done
		}
		// 5. Tell the gateway tunserver we are starting streaming
		err = gatewayStream.SendMsg(&rpc.StartStreaming{})
		if err != nil {
			f.mu.Unlock()
			if err == io.EOF { //nolint:errorlint
				var frame grpctool.RawFrame
				err = gatewayStream.RecvMsg(&frame) // get the real error
			}
			_ = f.rpcAPI.HandleIOError(log, "SendMsg(StartStreaming)", err)
			return nil, retry.Backoff
		}
		f.done = true
		f.mu.Unlock()
		f.pollCancel()
		rt := ReadyGateway{
			URL:          gatewayURL,
			Stream:       gatewayStream,
			Conn:         gatewayConn,
			StreamCancel: cancel,
		}
		select {
		case <-attemptCtx.Done():
		case f.foundGateway <- rt:
			success = true
		}
		return nil, retry.Done
	})
}

func (f *gatewayFinder) maybeStopTrying(tryingGatewayURL string) {
	if tryingGatewayURL == f.ownPrivateAPIURL {
		return // keep trying the own URL
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, gatewayURL := range f.gatewayURLs {
		if gatewayURL == tryingGatewayURL {
			return // known URLs still contain this URL so keep trying it.
		}
	}
	attempt := f.connections[tryingGatewayURL]
	delete(f.connections, tryingGatewayURL)
	attempt.cancel()
}

func (f *gatewayFinder) stopAllConnectionAttemptsExcept(gatewayURL string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	for url, c := range f.connections {
		if url != gatewayURL {
			c.cancel()
		}
	}
}

func (f *gatewayFinder) stopAllConnectionAttempts() {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, c := range f.connections {
		c.cancel()
	}
}

func stopAndDrain(t *time.Timer) {
	if !t.Stop() {
		select {
		case <-t.C:
		default:
		}
	}
}