package router

import (
	"context"

	"gitlab.com/gitlab-org/cluster-integration/tunnel/tunserver"
	otelcodes "go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type findHandle struct {
	tracer       trace.Tracer
	retTunHolder <-chan *tunnelHolder
	done         func()
	gotTunnel    bool
}

func (h *findHandle) Get(ctx context.Context) (tunserver.Tunnel, error) {
	ctx, span := h.tracer.Start(ctx, "findHandle.Get", trace.WithSpanKind(trace.SpanKindInternal))
	defer span.End()

	select {
	case <-ctx.Done():
		span.SetStatus(otelcodes.Error, "FindTunnel request aborted")
		span.RecordError(ctx.Err())
		return nil, ctx.Err()
	case tunHolder := <-h.retTunHolder:
		h.gotTunnel = true
		if tunHolder == nil {
			span.SetStatus(otelcodes.Error, "kas is shutting down")
			return nil, status.Error(codes.Canceled, "kas is shutting down")
		}
		span.SetStatus(otelcodes.Ok, "")
		return tunHolder, nil
	}
}

func (h *findHandle) Done(ctx context.Context) {
	_, span := h.tracer.Start(ctx, "findHandle.Done", trace.WithSpanKind(trace.SpanKindInternal))
	defer span.End()

	if h.gotTunnel {
		// No cleanup needed if Get returned a tunnel.
		return
	}
	h.done()
}
