package tunserver

import (
	"context"
	"io"
	"log/slog"

	"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/tunnel/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/mem"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

var (
	_ grpc.ClientStream = (*routingStream)(nil)
)

type routingStream struct {
	log               *slog.Logger
	api               modshared.API
	ro                routingOpts
	agentKey          api.AgentKey
	destination       grpc.ClientStream
	destinationCancel context.CancelFunc
	codec             encoding.CodecV2
	recvError         error
	header            metadata.MD
	trailer           metadata.MD
	haveHeader        bool
	haveRecvErr       bool
}

func (s *routingStream) Header() (metadata.MD, error) {
	if s.haveHeader {
		return s.header, nil
	}
	if s.haveRecvErr {
		// Do not return the error. The user should get it by calling RecvMsg().
		return nil, nil
	}
	tmp := &rpc.GatewayResponse{}
	err := s.destination.RecvMsg(tmp)
	if err != nil {
		s.recvError = err
		s.haveRecvErr = true
		// Do not return the error. The user should get it by calling RecvMsg().
		return nil, nil //nolint:nilerr
	}
	switch msg := tmp.Msg.(type) {
	case *rpc.GatewayResponse_Header_:
		s.setHeader(msg.Header.Metadata())
		return s.header, nil
	default:
		s.recvError = status.Errorf(codes.Internal, "unexpected message type: %T", tmp.Msg)
		s.haveRecvErr = true
		s.reportError(s.recvError)
		s.destinationCancel() // unblock SendMsg()
		// Do not return the error. The user should get it by calling RecvMsg().
		return nil, nil
	}
}

func (s *routingStream) Trailer() metadata.MD {
	return s.trailer
}

func (s *routingStream) CloseSend() error {
	return s.destination.CloseSend()
}

func (s *routingStream) Context() context.Context {
	return s.destination.Context()
}

func (s *routingStream) SendMsg(m any) error {
	return s.destination.SendMsg(m)
}

func (s *routingStream) RecvMsg(m any) error {
	if s.haveRecvErr {
		return s.recvError
	}
	for {
		tmp := &rpc.GatewayResponse{}
		err := s.destination.RecvMsg(tmp)
		if err != nil {
			s.recvError = err
			s.haveRecvErr = true
			return err
		}
		switch msg := tmp.Msg.(type) {
		case *rpc.GatewayResponse_Header_:
			if s.haveHeader {
				s.recvError = status.Errorf(codes.Internal, "unexpected message type: %T", tmp.Msg)
				s.haveRecvErr = true
				s.reportError(s.recvError)
				s.destinationCancel() // unblock SendMsg()
				return s.recvError
			}
			s.setHeader(msg.Header.Metadata())
		case *rpc.GatewayResponse_Message_:
			data := mem.BufferSlice{mem.SliceBuffer(msg.Message.Data)}
			err = s.codec.Unmarshal(data, m)
			if err != nil {
				s.recvError = status.Errorf(codes.Internal, "codec.Unmarshal(): %v", err)
				s.haveRecvErr = true
				s.reportError(s.recvError)
				s.destinationCancel() // unblock SendMsg()
				return s.recvError
			}
			return nil
		case *rpc.GatewayResponse_Trailer_:
			s.trailer = msg.Trailer.Metadata()
			for _, t := range s.ro.trailer {
				*t = s.trailer
			}
		case *rpc.GatewayResponse_Error_:
			s.recvError = status.ErrorProto(msg.Error.Status)
			s.haveRecvErr = true
			err = s.destination.RecvMsg(tmp) // consume io.EOF
			switch err {
			case io.EOF: //nolint:errorlint
				// as expected
			case nil:
				// This is a protocol violation, report that.
				s.recvError = status.Errorf(codes.Internal, "unexpected message after GatewayResponse_Error: %T. Error was: %v", tmp.Msg, s.recvError)
				s.reportError(s.recvError)
				s.destinationCancel() // unblock SendMsg()
			default:
				// Report the last error from the stream - we failed to finish reading the stream,
				// so we cannot say it was a "success".
				s.recvError = err
			}
			return s.recvError
		default:
			s.recvError = status.Errorf(codes.Internal, "unexpected message type: %T", tmp.Msg)
			s.haveRecvErr = true
			s.reportError(s.recvError)
			s.destinationCancel() // unblock SendMsg()
			return s.recvError
		}
	}
}

func (s *routingStream) setHeader(m metadata.MD) {
	s.header = m
	s.haveHeader = true
	for _, h := range s.ro.header {
		*h = m
	}
}

func (s *routingStream) reportError(err error) {
	s.api.HandleProcessingError(context.Background(), s.log, "Routing error", err, fieldz.AgentKey(s.agentKey))
}
