main.go

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	gcppubsub "cloud.google.com/go/pubsub/v2"
	"github.com/dipjyotimetia/pubsub-emulator/internal/config"
	"github.com/dipjyotimetia/pubsub-emulator/internal/dashboard"
	"github.com/dipjyotimetia/pubsub-emulator/internal/pubsub"
	"github.com/dipjyotimetia/pubsub-emulator/internal/server"
	"github.com/dipjyotimetia/pubsub-emulator/pkg/logger"
)

const (
	// startupAckDeadlineSeconds is the ack deadline applied to subscriptions
	// created from configuration at startup. (Subscriptions created later via
	// the dashboard use their own default; see the dashboard package.)
	startupAckDeadlineSeconds = 20
	// subscriberDrainTimeout bounds how long shutdown waits for in-flight
	// message receivers to stop after the context is cancelled.
	subscriberDrainTimeout = 10 * time.Second
)

func main() {
	// Initialize logger
	log := logger.New()
	log.Info("Starting Pub/Sub Emulator with refactored architecture")

	// Load configuration
	cfg, err := config.LoadFromEnv()
	if err != nil {
		log.Fatal("Failed to load configuration: %v", err)
	}

	log.Info("Configuration loaded: Project=%s, Topics=%v, Subscriptions=%v",
		cfg.ProjectID, cfg.TopicIDs, cfg.SubscriptionIDs)

	// Context cancelled on SIGINT/SIGTERM; drives both the subscribers and the
	// HTTP server so shutdown propagates everywhere.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	// Initialize Pub/Sub client wrapper (creates GCP client internally)
	psClient, err := pubsub.NewClient(ctx, cfg.ProjectID, log)
	if err != nil {
		log.Fatal("Failed to create Pub/Sub client: %v", err)
	}
	defer func() { _ = psClient.Close() }()

	// Get underlying GCP client for dashboard
	pubsubClient := psClient.GetClient()

	// Create topics and subscriptions
	if err := setupTopicsAndSubscriptions(ctx, psClient, cfg, log); err != nil {
		log.Fatal("Failed to setup topics and subscriptions: %v", err)
	}

	// Initialize dashboard
	dash := dashboard.New(pubsubClient, cfg.ProjectID, log)

	// Initialize publisher
	pub := pubsub.NewPublisher(psClient, log)

	// Publish initial messages to topics
	if err := publishInitialMessages(ctx, pub, cfg, dash, log); err != nil {
		log.Error("Failed to publish initial messages: %v", err)
	}

	// Initialize subscriber and start receiving messages from subscriptions
	sub := pubsub.NewSubscriber(psClient, log)
	wg := startSubscriptions(ctx, sub, cfg, dash, log)

	// Initialize and start HTTP server with graceful shutdown
	srv := server.New(&server.Config{
		Port:      cfg.DashboardPort,
		Dashboard: dash,
		Logger:    log,
	})

	log.Info("Dashboard will be available at http://localhost:%s", cfg.DashboardPort)

	// Blocks until ctx is cancelled (signal) or the server fails.
	serverErr := srv.Start(ctx)

	// Cancel the context explicitly so subscribers also stop on the
	// server-error path (where no signal cancelled it), then drain receivers.
	stop()
	waitForSubscribers(wg, log)

	if serverErr != nil {
		log.Error("Server error: %v", serverErr)
	}

	log.Info("Application shutdown complete")
}

// waitForSubscribers waits for all subscription receivers to stop, bounded by
// subscriberDrainTimeout so shutdown can never hang indefinitely.
func waitForSubscribers(wg *sync.WaitGroup, log *logger.Logger) {
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		log.Info("All subscription receivers stopped")
	case <-time.After(subscriberDrainTimeout):
		log.Warn("Timed out waiting for subscription receivers to drain")
	}
}

// setupTopicsAndSubscriptions creates topics and subscriptions
func setupTopicsAndSubscriptions(ctx context.Context, psClient *pubsub.Client, cfg *config.Config, log *logger.Logger) error {
	if len(cfg.TopicIDs) != len(cfg.SubscriptionIDs) {
		return fmt.Errorf("number of topics (%d) and subscriptions (%d) must match",
			len(cfg.TopicIDs), len(cfg.SubscriptionIDs))
	}

	for i, topicID := range cfg.TopicIDs {
		// Create topic
		topic, err := psClient.CreateTopic(ctx, topicID)
		if err != nil {
			log.Warn("Failed to create topic %s: %v (may already exist)", topicID, err)
		} else {
			log.Info("Created topic: %s", topic.Name)
		}

		// Create subscription
		subID := cfg.SubscriptionIDs[i]
		subscription, err := psClient.CreateSubscription(ctx, subID, topicID, startupAckDeadlineSeconds)
		if err != nil {
			log.Warn("Failed to create subscription %s: %v (may already exist)", subID, err)
		} else {
			log.Info("Created subscription: %s", subscription.Name)
		}
	}

	return nil
}

// publishInitialMessages publishes messages to all configured topics
func publishInitialMessages(ctx context.Context, pub *pubsub.Publisher, cfg *config.Config, dash *dashboard.Dashboard, log *logger.Logger) error {
	message := cfg.MessageToPublish
	if message == "" {
		message = "Hello from Pub/Sub Emulator!"
	}

	attributes := map[string]string{
		"source": "emulator",
		"time":   time.Now().Format(time.RFC3339),
	}

	log.Info("Publishing initial message to %d topics", len(cfg.TopicIDs))

	for _, topicID := range cfg.TopicIDs {
		msgID, err := pub.PublishMessage(ctx, topicID, message, attributes)
		if err != nil {
			log.Error("Failed to publish to topic %s: %v", topicID, err)
			continue
		}

		log.Info("Published message to %s with ID: %s", topicID, msgID)

		// Add to dashboard
		msg := &gcppubsub.Message{
			ID:          msgID,
			Data:        []byte(message),
			Attributes:  attributes,
			PublishTime: time.Now(),
		}
		dash.AddMessage(msg, topicID)
	}

	return nil
}

// startSubscriptions starts message receivers for all subscriptions and returns
// a WaitGroup that completes once the receivers stop.
func startSubscriptions(ctx context.Context, sub *pubsub.Subscriber, cfg *config.Config, dash *dashboard.Dashboard, log *logger.Logger) *sync.WaitGroup {
	log.Info("Starting message receivers for %d subscriptions", len(cfg.SubscriptionIDs))

	// Create subscription to topic mapping
	subToTopicMap := make(map[string]string)
	for i := range cfg.SubscriptionIDs {
		if i < len(cfg.TopicIDs) {
			subToTopicMap[cfg.SubscriptionIDs[i]] = cfg.TopicIDs[i]
		}
	}

	// Create handler function that adds messages to dashboard
	handler := func(ctx context.Context, msg *gcppubsub.Message) {
		// Extract subscription ID from message (if available)
		// Note: This is a simplified approach. In production, you'd need to
		// track which subscription received the message
		topicID := ""
		if len(cfg.TopicIDs) > 0 {
			topicID = cfg.TopicIDs[0] // Fallback to first topic
		}

		// Try to get topic from message attributes if available
		if topic, ok := msg.Attributes["_topic"]; ok {
			topicID = topic
		}

		log.Debug("Received message: %s from topic: %s", msg.ID, topicID)
		dash.AddMessage(msg, topicID)
	}

	// Subscribe to all subscriptions
	return sub.SubscribeToAll(ctx, cfg.SubscriptionIDs, cfg.TopicIDs, handler)
}