package router

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/tunserver"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type findTunnelRequest struct {
	service, method string
	retTun          chan<- *tunserver.TunnelImpl
}

type findHandle struct {
	retTun    <-chan *tunserver.TunnelImpl
	done      func(context.Context)
	gotTunnel bool
}

func (h *findHandle) Get(ctx context.Context) (tunserver.Tunnel, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case tun := <-h.retTun:
		h.gotTunnel = true
		if tun == nil {
			return nil, status.Error(codes.Canceled, "agentk is shutting down")
		}
		return tun, nil
	}
}

func (h *findHandle) Done(ctx context.Context) {
	if h.gotTunnel {
		// No cleanup needed if Get returned a tunnel.
		return
	}
	h.done(ctx)
}

type Registry struct {
	log *slog.Logger

	mu           sync.Mutex
	findRequests map[*findTunnelRequest]struct{}
	tuns         map[*tunserver.TunnelImpl]struct{}
}

func NewRegistry(log *slog.Logger) *Registry {
	return &Registry{
		log:          log,
		findRequests: make(map[*findTunnelRequest]struct{}),
		tuns:         make(map[*tunserver.TunnelImpl]struct{}),
	}
}

func (r *Registry) FindTunnel(_ context.Context, service, method string) (bool, tunserver.FindHandle) {
	// Buffer 1 to not block on send when a tunnel is found before find request is registered.
	retTun := make(chan *tunserver.TunnelImpl, 1) // can receive nil from it if Stop() is called
	ftr := &findTunnelRequest{
		service: service,
		method:  method,
		retTun:  retTun,
	}
	found := false
	func() {
		r.mu.Lock()
		defer r.mu.Unlock()

		// 1. Check if we have a suitable tunnel
		for tun := range r.tuns {
			if !tun.Descriptor.SupportsServiceAndMethod(service, method) {
				continue
			}
			// Suitable tunnel found!
			tun.State = tunserver.StateFound
			retTun <- tun // must not block because the reception is below
			found = true
			r.unregisterTunnelLocked(tun)
			return
		}
		// 2. No suitable tunnel found, add to the queue
		r.findRequests[ftr] = struct{}{}
	}()
	return found, &findHandle{
		retTun: retTun,
		done: func(ctx context.Context) {
			r.mu.Lock()
			defer r.mu.Unlock()
			close(retTun)
			tun := <-retTun // will get nil if there was nothing in the channel or if registry is shutting down.
			if tun != nil {
				// Got the tunnel, but it's too late so return it to the registry.
				r.onTunnelDoneLocked(tun)
			} else {
				r.deleteFindRequestLocked(ftr)
			}
		},
	}
}

// HandleTunnel is called with server-side interface of the reverse tunnel.
// It registers the tunnel and blocks, waiting for a request to proxy through the tunnel.
// The method returns the error value to return to gRPC framework.
// ageCtx can be used to unblock the method if the tunnel is not being used already.
func (r *Registry) HandleTunnel(ageCtx context.Context, server grpc.BidiStreamingServer[rpc.ConnectRequest, rpc.ConnectResponse]) error {
	recv, err := server.Recv()
	if err != nil {
		return err
	}
	descriptor, ok := recv.Msg.(*rpc.ConnectRequest_Descriptor_)
	if !ok {
		return status.Errorf(codes.InvalidArgument, "invalid oneof value type: %T", recv.Msg)
	}
	retErr := make(chan error, 1)
	tun := &tunserver.TunnelImpl{
		Tunnel:       server,
		TunnelRetErr: retErr,
		AgentKey:     modshared.NoAgentKey,
		Descriptor:   descriptor.Descriptor_.ApiDescriptor,
		State:        tunserver.StateReady,
		OnForward:    r.onTunnelForward,
		OnDone:       r.onTunnelDone,
	}
	// Register
	r.registerTunnel(tun)
	// Wait for return error or for cancellation
	select {
	case <-ageCtx.Done():
		// Context canceled
		r.mu.Lock()
		switch tun.State {
		case tunserver.StateReady:
			tun.State = tunserver.StateContextDone
			r.unregisterTunnelLocked(tun)
			r.mu.Unlock()
			return nil
		case tunserver.StateFound:
			// Tunnel was found but hasn't been used yet, Done() hasn't been called.
			// Set State to StateContextDone so that ForwardStream() errors out without doing any I/O.
			tun.State = tunserver.StateContextDone
			r.mu.Unlock()
			return nil
		case tunserver.StateForwarding:
			// I/O on the stream will error out, just wait for the return value.
			r.mu.Unlock()
			return <-retErr
		case tunserver.StateDone:
			// Forwarding has finished and then ctx signaled done. Return the result value from forwarding.
			r.mu.Unlock()
			return <-retErr
		case tunserver.StateContextDone:
			// Cannot happen twice.
			r.mu.Unlock()
			panic(errors.New("unreachable"))
		default:
			// Should never happen
			r.mu.Unlock()
			panic(fmt.Errorf("invalid State: %d", tun.State))
		}
	case err = <-retErr:
		return err
	}
}

func (r *Registry) registerTunnel(toReg *tunserver.TunnelImpl) {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.registerTunnelLocked(toReg)
}

func (r *Registry) registerTunnelLocked(toReg *tunserver.TunnelImpl) {
	// 1. Before registering the tunnel see if there is a find tunnel request waiting for it
	for ftr := range r.findRequests {
		if !toReg.Descriptor.SupportsServiceAndMethod(ftr.service, ftr.method) {
			continue
		}
		// Waiting request found!
		toReg.State = tunserver.StateFound

		ftr.retTun <- toReg            // Satisfy the waiting request ASAP
		r.deleteFindRequestLocked(ftr) // Remove it from the queue
		return
	}

	// 2. Register the tunnel
	toReg.State = tunserver.StateReady
	r.tuns[toReg] = struct{}{}
}

func (r *Registry) unregisterTunnelLocked(toUnreg *tunserver.TunnelImpl) {
	delete(r.tuns, toUnreg)
}

func (r *Registry) onTunnelForward(tun *tunserver.TunnelImpl) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	switch tun.State {
	case tunserver.StateReady:
		return status.Error(codes.Internal, "unreachable: ready -> forwarding should never happen")
	case tunserver.StateFound:
		tun.State = tunserver.StateForwarding
		return nil
	case tunserver.StateForwarding:
		return status.Error(codes.Internal, "ForwardStream() called more than once")
	case tunserver.StateDone:
		return status.Error(codes.Internal, "ForwardStream() called after Done()")
	case tunserver.StateContextDone:
		return status.Error(codes.Canceled, "ForwardStream() called on done stream")
	default:
		return status.Errorf(codes.Internal, "unreachable: invalid State: %d", tun.State)
	}
}

func (r *Registry) onTunnelDone(_ context.Context, tun *tunserver.TunnelImpl) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.onTunnelDoneLocked(tun)
}

func (r *Registry) onTunnelDoneLocked(tun *tunserver.TunnelImpl) {
	switch tun.State {
	case tunserver.StateReady:
		panic(errors.New("unreachable: ready -> done should never happen"))
	case tunserver.StateFound:
		// Tunnel was found but was not used, Done() was called. Just put it back.
		r.registerTunnelLocked(tun)
	case tunserver.StateForwarding:
		tun.State = tunserver.StateDone
	case tunserver.StateDone:
		panic("Done() called more than once")
	case tunserver.StateContextDone:
	// Done() called after canceled context in HandleTunnel(). Nothing to do.
	default:
		// Should never happen
		panic(fmt.Sprintf("invalid State: %d", tun.State))
	}
}

func (r *Registry) deleteFindRequestLocked(ftr *findTunnelRequest) {
	delete(r.findRequests, ftr)
}
