package grpcserver

import (
	"log"

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/proxy"
	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/services"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// ProxyStream carries one console TCP connection. The agent opens it after
// dialling its own loopback address, and authenticates with the same agent
// token as the command stream plus the single-use proxy_id it was handed.
func (s *vantageServer) ProxyStream(stream pb.Vantage_ProxyStreamServer) error {
	msg, err := stream.Recv()
	if err != nil {
		return status.Errorf(codes.InvalidArgument, "expected initial open message: %v", err)
	}
	if msg.Open == nil {
		return status.Error(codes.InvalidArgument, "first message must be open")
	}

	srv, err := services.ValidateAgentToken(msg.Open.ServerId, msg.Open.AgentToken)
	if err != nil {
		// Deliberately identical to the claim-failure response below: a bad
		// token, an unknown proxy_id, and a proxy_id belonging to another
		// instance/server must be indistinguishable to the caller.
		log.Printf("proxy %s (server %s): invalid agent token", msg.Open.ProxyId, msg.Open.ServerId)
		return status.Error(codes.PermissionDenied, "proxy session unavailable")
	}

	if err := serveProxy(msg.Open, srv.InstanceID, stream); err != nil {
		// The reason is deliberately not returned to the agent: an unknown and a
		// foreign proxy_id must be indistinguishable.
		log.Printf("proxy %s (server %s): %v", msg.Open.ProxyId, msg.Open.ServerId, err)
		return status.Error(codes.PermissionDenied, "proxy session unavailable")
	}
	return nil
}

// serveProxy claims the pending session, binds this pod's relay listener for
// it, and relays. Split out from the gRPC method so the authorisation matrix is
// testable without a real stream.
//
// The listener is bound here, on whichever replica the stream reached, rather
// than in advance on the pod holding the agent's command stream — those are not
// the same pod, because an L7 proxy balances HTTP/2 requests independently.
func serveProxy(open *pb.ProxyOpen, instanceID string, stream proxy.AgentStream) error {
	sess, err := services.ClaimProxyStream(instanceID, open.ServerId, open.ProxyId)
	if err != nil {
		return err
	}
	return sess.Serve(stream)
}
