Turbine/pkg/ring/main.go

package main

import (
	"fmt"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"git.solsynth.dev/goatworks/turbine/pkg/ring/clients" // Add this import
	"git.solsynth.dev/goatworks/turbine/pkg/ring/infra"
	"git.solsynth.dev/goatworks/turbine/pkg/ring/routes"
	"git.solsynth.dev/goatworks/turbine/pkg/ring/services"
	"git.solsynth.dev/goatworks/turbine/pkg/ring/websocket" // Add this import
	"git.solsynth.dev/goatworks/turbine/pkg/shared/hash"
	pb "git.solsynth.dev/goatworks/turbine/pkg/shared/proto/gen"
	"git.solsynth.dev/goatworks/turbine/pkg/shared/registrar"

	fiber_websocket "github.com/gofiber/contrib/v3/websocket" // Add this alias import
	"github.com/gofiber/fiber/v3"
	"github.com/google/uuid"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/samber/lo"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
)

func init() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
}

func main() {
	log.Info().Msg("Starting Turbine Ring...")
	viper.SetConfigName("settings")
	viper.AddConfigPath(".")
	viper.AddConfigPath("/etc/turbine/")
	viper.SetConfigType("toml")

	log.Info().Msg("Reading configuration...")
	if err := viper.ReadInConfig(); err != nil {
		log.Fatal().Err(err).Msg("Failed to read config file...")
	}
	log.Info().Msg("Configuration loaded.")

	clients.InitPushClients()  // Initialize APNs and Firebase clients
	clients.InitSMTPSettings() // Initialize SMTP client settings
	clients.InitNATSClient()   // Initialize NATS client

	// --- gRPC Server ---
	grpcListenAddr := viper.GetString("grpc.listen")
	if grpcListenAddr == "" {
		log.Fatal().Msg("grpc.listen not configured")
	}

	grpcServer := grpc.NewServer()
	wsManager := websocket.NewManager(clients.GetNATSClient()) // Pass NATS client to WebSocket Manager
	ringService := &services.RingServiceServerImpl{
		WsManager: wsManager, // Inject WebSocket Manager into RingService
	}
	pb.RegisterRingServiceServer(grpcServer, ringService)
	pb.RegisterRingHandlerServiceServer(grpcServer, ringService)

	grpcLis, err := net.Listen("tcp", grpcListenAddr)
	if err != nil {
		log.Fatal().Err(err).Msgf("Failed to listen for gRPC: %s", grpcListenAddr)
	}

	go func() {
		log.Info().Msgf("gRPC server listening on %s", grpcListenAddr)
		if err := grpcServer.Serve(grpcLis); err != nil {
			log.Fatal().Err(err).Msg("Failed to serve gRPC")
		}
	}()

	// --- Service Registration ---
	etcdEndpoints := viper.GetStringSlice("etcd.endpoints")
	if len(etcdEndpoints) == 0 {
		log.Fatal().Msg("etcd.endpoints not configured")
	}

	if err := infra.ConnectDb(); err != nil {
		log.Fatal().Err(err).Msg("Failed to connect database.")
	} else {
		infra.Db.AutoMigrate()
	}

	if viper.GetBool("etcd.insecure") {
		for i, ep := range etcdEndpoints {
			if !strings.HasPrefix(ep, "http://") && !strings.HasPrefix(ep, "https://") {
				etcdEndpoints[i] = "http://" + ep
			}
		}
	}

	serviceReg, err := registrar.NewServiceRegistrar(etcdEndpoints)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to create service registrar")
	}

	listenAddr := viper.GetString("listen")
	host := viper.GetString("host")
	if host == "" {
		log.Fatal().Msg("host not configured")
	}
	portStr := strings.TrimPrefix(listenAddr, ":")
	port, err := strconv.Atoi(portStr)
	if err != nil {
		log.Fatal().Err(err).Msg("Invalid listen address")
	}

	serviceName := "ring" // This should probably be "ring"
	instanceID := fmt.Sprint(hash.Hash(fmt.Sprintf("%s-%s-%d", serviceName, host, port)))[:8]

	err = serviceReg.Register(serviceName, "http", instanceID, host, port, 30)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to register service")
	}
	log.Info().Str("service", serviceName).Str("instanceID", instanceID).Msg("Service registered successfully")

	// --- Web Server ---
	app := fiber.New(fiber.Config{
		ServerHeader: "Turbine Ring",
	})

	api := app.Group("/api")
	{
		api.Post("/notifications/subscription", routes.CreatePushSubscription)
	}

	// Initialize WebSocket Controller
	wsController := websocket.NewWebSocketController(wsManager)

	// WebSocket endpoint
	app.Use("/ws", func(c fiber.Ctx) error {
		// Mock authentication based on C# example
		// In a real scenario, you'd extract user/session from JWT or similar
		// and set c.Locals("currentUser") and c.Locals("currentSession")
		c.Locals("currentUser", &pb.Account{Id: uuid.New().String(), Name: "mock_user", Nick: "Mock User"})
		c.Locals("currentSession", &pb.AuthSession{ClientId: lo.ToPtr(uuid.New().String())})

		if fiber_websocket.IsWebSocketUpgrade(c) {
			return c.Next()
		}
		return fiber.ErrUpgradeRequired
	})
	app.Get("/ws", fiber_websocket.New(wsController.HandleWebSocket))

	// Graceful shutdown
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-c
		log.Info().Msg("Shutting down servers...")
		if err := app.ShutdownWithTimeout(5 * time.Second); err != nil {
			log.Error().Err(err).Msg("Fiber server shutdown error")
		}
		grpcServer.GracefulStop()
		log.Info().Msg("Servers gracefully stopped")
	}()

	err = app.Listen(listenAddr, fiber.ListenConfig{DisableStartupMessage: true})
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to start the server...")
	}
}