package tunserver

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"slices"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/fieldz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/retry"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/rpc"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/encoding"
	"google.golang.org/protobuf/reflect/protoreflect"
	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	tunnelReadyFieldNumber protoreflect.FieldNumber = 1
	headerFieldNumber      protoreflect.FieldNumber = 2
	messageFieldNumber     protoreflect.FieldNumber = 3
	trailerFieldNumber     protoreflect.FieldNumber = 4
	errorFieldNumber       protoreflect.FieldNumber = 5
	noTunnelFieldNumber    protoreflect.FieldNumber = 6
)

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

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

type Target interface {
	comparable
	String() string
}

type connAttempt struct {
	cancel context.CancelFunc
}

type ReadyGateway[T Target] struct {
	URL          T
	Stream       grpc.ClientStream
	StreamCancel context.CancelFunc
	Codec        encoding.CodecV2
}

// PollGatewayURLsCallback is called periodically with found gateway tunserver URLs for a particular agent key.
type PollGatewayURLsCallback[T Target] func(gatewayURLs []T)

type PollingGatewayURLQuerier[T Target] interface {
	PollGatewayURLs(ctx context.Context, agentKey api.AgentKey, cb PollGatewayURLsCallback[T])
	CachedGatewayURLs(agentKey api.AgentKey) []T
}

type GatewayFinder[T Target] interface {
	Find(ctx context.Context) (ReadyGateway[T], error)
}

type gatewayFinder[T Target] struct {
	outgoingCtx           context.Context //nolint:containedctx
	log                   *slog.Logger
	gatewayPool           grpctool.PoolInterface[T]
	gatewayQuerier        PollingGatewayURLQuerier[T]
	api                   modshared.API
	fullMethod            string // /service/method
	ownPrivateAPIURL      T
	agentKey              api.AgentKey
	pollConfig            retry.PollConfigFactory
	foundGateway          chan ReadyGateway[T]
	noTunnel              chan struct{}
	wg                    wait.Group
	pollCancel            context.CancelFunc
	tryNewGatewayInterval time.Duration

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

func NewGatewayFinder[T Target](outgoingCtx context.Context, log *slog.Logger, gatewayPool grpctool.PoolInterface[T],
	gatewayQuerier PollingGatewayURLQuerier[T], modAPI modshared.API, fullMethod string,
	ownPrivateAPIURL T, agentKey api.AgentKey, pollConfig retry.PollConfigFactory, tryNewGatewayInterval time.Duration) GatewayFinder[T] {
	return &gatewayFinder[T]{
		outgoingCtx:           outgoingCtx,
		log:                   log,
		gatewayPool:           gatewayPool,
		gatewayQuerier:        gatewayQuerier,
		api:                   modAPI,
		fullMethod:            fullMethod,
		ownPrivateAPIURL:      ownPrivateAPIURL,
		agentKey:              agentKey,
		pollConfig:            pollConfig,
		tryNewGatewayInterval: tryNewGatewayInterval,
		foundGateway:          make(chan ReadyGateway[T]),
		noTunnel:              make(chan struct{}),
		connections:           make(map[T]connAttempt),
	}
}

func (f *gatewayFinder[T]) Find(ctx context.Context) (ReadyGateway[T], 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 []T)
	f.gatewayURLs = f.gatewayQuerier.CachedGatewayURLs(f.agentKey)
	done := ctx.Done()

	// Timer must have been stopped or has fired when this function is called
	tryNewGatewayWhenTimerNotRunning := func() {
		if f.tryNewGateway() {
			// 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.agentKey, func(gatewayURLs []T) {
						select {
						case <-pollDone:
						case gatewayURLsC <- gatewayURLs:
						}
					})
				})
			}
		}
	}

	for {
		select {
		case <-done:
			f.stopAllConnectionAttempts()
			return ReadyGateway[T]{}, 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() {
				// 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[T]) 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[T]) tryGatewayLocked(gatewayURL T) {
	connCtx, connCancel := context.WithCancel(f.outgoingCtx)
	f.connections[gatewayURL] = connAttempt{
		cancel: connCancel,
	}
	f.wg.Start(func() {
		f.tryGatewayAsync(connCtx, connCancel, gatewayURL)
	})
}

func (f *gatewayFinder[T]) tryGatewayAsync(ctx context.Context, cancel context.CancelFunc, gatewayURL T) {
	log := f.log.With(logz.GatewayURL(gatewayURL.String()))
	noTunnelSent := false
	_ = retry.PollWithBackoff(ctx, f.pollConfig(), func(ctx context.Context) (error, retry.AttemptResult) { //nolint:staticcheck
		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(gatewayURL)
		if err != nil {
			f.api.HandleProcessingError(attemptCtx, log, "Failed to dial gateway tunserver", err, fieldz.AgentKey(f.agentKey))
			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.WaitForReady(true),
		)
		if err != nil {
			f.api.HandleProcessingError(attemptCtx, log, "Failed to open a new stream to gateway tunserver", err, fieldz.AgentKey(f.agentKey))
			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 errTunnelReadySentinel
			}),
			grpctool.WithNotExpectingToGet(codes.Internal, headerFieldNumber, messageFieldNumber, trailerFieldNumber, errorFieldNumber),
		)
		switch err {
		case nil:
			// Gateway tunserver closed the connection cleanly.
			// Perhaps it's been open for too long or the server is shutting down.
			return nil, retry.ContinueImmediately
		case errTunnelReadySentinel: //nolint:errorlint
			// fallthrough
		default:
			f.api.HandleProcessingError(attemptCtx, log, "RecvMsg(GatewayResponse)", err, fieldz.AgentKey(f.agentKey))
			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.api.HandleProcessingError(attemptCtx, log, "SendMsg(StartStreaming)", err, fieldz.AgentKey(f.agentKey))
			return nil, retry.Backoff
		}
		f.done = true
		f.mu.Unlock()
		f.pollCancel()
		rt := ReadyGateway[T]{
			URL:          gatewayURL,
			Stream:       gatewayStream,
			StreamCancel: cancel,
			Codec:        grpctool.ProtoCodec,
		}
		select {
		case <-attemptCtx.Done():
		case f.foundGateway <- rt:
			success = true
			context.AfterFunc(attemptCtx, gatewayConn.Done)
		}
		return nil, retry.Done
	})
}

func (f *gatewayFinder[T]) maybeStopTrying(tryingGatewayURL T) {
	if tryingGatewayURL == f.ownPrivateAPIURL {
		return // keep trying the own URL
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if slices.Contains(f.gatewayURLs, 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[T]) stopAllConnectionAttemptsExcept(gatewayURL T) {
	f.mu.Lock()
	defer f.mu.Unlock()
	for url, c := range f.connections {
		if url != gatewayURL {
			c.cancel()
		}
	}
}

func (f *gatewayFinder[T]) 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:
		}
	}
}
