package rpc

import (
	"net/http"

	"connectrpc.com/grpcreflect"
)

// ServeMux wraps http.ServeMux to collect registered Connect-RPC service names for automatic reflection.
type ServeMux struct {
	*http.ServeMux
	services []string
}

// NewServeMux creates a new initialized ServeMux.
func NewServeMux() *ServeMux {
	return &ServeMux{
		ServeMux: http.NewServeMux(),
	}
}

// RegisterConnect registers a Connect-RPC service handler and tracks the service name for reflection.
func (sm *ServeMux) RegisterConnect(serviceName string, path string, handler http.Handler) {
	sm.ServeMux.Handle(path, handler)
	sm.services = append(sm.services, serviceName)
}

// MountReflection registers the reflection handlers for all registered Connect services on the ServeMux.
func (sm *ServeMux) MountReflection() {
	if len(sm.services) == 0 {
		return
	}
	reflector := grpcreflect.NewStaticReflector(sm.services...)
	sm.ServeMux.Handle(grpcreflect.NewHandlerV1(reflector))
	sm.ServeMux.Handle(grpcreflect.NewHandlerV1Alpha(reflector))
}