package mtls

import (
	"context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"
)

// nodeUUIDKey is the unexported context key for the verified node UUID.
type nodeUUIDKey struct{}

// enrollFullMethod is the one RPC exempt from mTLS enforcement.
// The agent has no certificate yet when it calls Enroll; the bootstrap token
// provides authentication for that call instead.
const enrollFullMethod = "/pangolin.agent.v1.AgentService/Enroll"

// UnaryServerInterceptor returns a gRPC unary interceptor that enforces mTLS identity.
//
// For the Enroll method: the call is passed through without a certificate check.
// For all other methods: a verified client certificate is required; its CN is
// injected into the context as the authoritative node UUID.
func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(
		ctx context.Context,
		req interface{},
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler,
	) (interface{}, error) {
		if info.FullMethod == enrollFullMethod {
			return handler(ctx, req)
		}

		nodeUUID, err := extractNodeUUID(ctx, info.FullMethod)
		if err != nil {
			return nil, err
		}

		return handler(context.WithValue(ctx, nodeUUIDKey{}, nodeUUID), req)
	}
}

// StreamServerInterceptor returns a gRPC stream interceptor that enforces mTLS identity.
// Same whitelist logic as the unary interceptor applies.
func StreamServerInterceptor() grpc.StreamServerInterceptor {
	return func(
		srv interface{},
		ss grpc.ServerStream,
		info *grpc.StreamServerInfo,
		handler grpc.StreamHandler,
	) error {
		if info.FullMethod == enrollFullMethod {
			return handler(srv, ss)
		}

		nodeUUID, err := extractNodeUUID(ss.Context(), info.FullMethod)
		if err != nil {
			return err
		}

		enriched := &wrappedStream{
			ServerStream: ss,
			ctx:          context.WithValue(ss.Context(), nodeUUIDKey{}, nodeUUID),
		}
		return handler(srv, enriched)
	}
}

// NodeUUIDFromContext retrieves the verified node UUID that was injected by the
// interceptor. Returns ("", false) when the context carries no UUID (e.g. in
// Enroll handlers where the identity is not yet established).
func NodeUUIDFromContext(ctx context.Context) (string, bool) {
	v, ok := ctx.Value(nodeUUIDKey{}).(string)
	return v, ok && v != ""
}

// ─── internals ───────────────────────────────────────────────────────────────

// extractNodeUUID reads the verified TLS peer certificate's CN from the context.
// It returns an Unauthenticated status error when no valid cert is present.
func extractNodeUUID(ctx context.Context, method string) (string, error) {
	p, ok := peer.FromContext(ctx)
	if !ok {
		return "", status.Errorf(codes.Unauthenticated,
			"mtls: no peer info for method %s", method)
	}

	tlsInfo, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok {
		return "", status.Errorf(codes.Unauthenticated,
			"mtls: non-TLS connection for method %s", method)
	}

	chains := tlsInfo.State.VerifiedChains
	if len(chains) == 0 || len(chains[0]) == 0 {
		return "", status.Errorf(codes.Unauthenticated,
			"mtls: client certificate required for method %s", method)
	}

	cn := chains[0][0].Subject.CommonName
	if cn == "" {
		return "", status.Errorf(codes.Unauthenticated,
			"mtls: empty CN in client certificate for method %s", method)
	}
	return cn, nil
}

// wrappedStream overrides the context on a ServerStream.
type wrappedStream struct {
	grpc.ServerStream
	ctx context.Context
}

func (w *wrappedStream) Context() context.Context {
	return w.ctx
}