package main

import (
	"fmt"
	"sync"
	"time"
)

// Broker manages subscriptions and publishes messages.
type Broker struct {
	mu          sync.RWMutex
	subscribers map[string][]chan string // Topic -> list of subscriber channels
}

func NewBroker() *Broker {
	return &Broker{
		subscribers: make(map[string][]chan string),
	}
}

// Subscribe registers a new subscriber for a given topic.
// Returns a channel on which messages for the topic will be received.
func (b *Broker) Subscribe(topic string) <-chan string {
	msgCh := make(chan string, 1)

	b.mu.Lock()
	defer b.mu.Unlock()

	b.subscribers[topic] = append(b.subscribers[topic], msgCh)
	fmt.Printf("Subscriber joined topic '%s'. Total subscribers: %d
", topic, len(b.subscribers[topic]))
	return msgCh
}

// Publish sends a message to all subscribers of a given topic.
func (b *Broker) Publish(topic, message string) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	subs := b.subscribers[topic]
	fmt.Printf("Publisher: Publishing '%s' to topic '%s' for %d subscribers.
", message, topic, len(subs))
	for _, subCh := range subs {
		// Use select with default to avoid blocking if a subscriber is slow
		select {
		case subCh <- message:
			// Message sent successfully
		default:
			fmt.Printf("Warning: Subscriber on topic '%s' is full, dropping message: %s
", topic, message)
		}
	}
}

// Unsubscribe removes a subscriber from a topic and closes its channel.
func (b *Broker) Unsubscribe(topic string, subCh <-chan string) {
	b.mu.Lock()
	defer b.mu.Unlock()

	subs := b.subscribers[topic]
	for i, ch := range subs {
		if ch == subCh {
			b.subscribers[topic] = append(subs[:i], subs[i+1:]...)
			close(ch) // Close the subscriber's channel
			fmt.Printf("Subscriber left topic '%s'. Remaining: %d
", topic, len(b.subscribers[topic]))
			return
		}
	}
}

func main() {
	broker := NewBroker()

	// Subscriber 1 for "news"
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		newsSubCh := broker.Subscribe("news")
		for msg := range newsSubCh {
			fmt.Printf("News Subscriber 1: %s
", msg)
			if msg == "The End" {
				fmt.Println("News Subscriber 1: Done receiving.")
				return
			}
		}
		fmt.Println("News Subscriber 1: Channel closed.")
	}()

	// Subscriber 2 for "news" and "alerts"
	wg.Add(1)
	go func() {
		defer wg.Done()
		newsSubCh := broker.Subscribe("news")
		alertSubCh := broker.Subscribe("alerts")
		for {
			select {
			case msg, ok := <-newsSubCh:
				if !ok { fmt.Println("News Subscriber 2 (news): Channel closed."); newsSubCh = nil }
				else { fmt.Printf("News Subscriber 2: %s
", msg) }
			case msg, ok := <-alertSubCh:
				if !ok { fmt.Println("News Subscriber 2 (alerts): Channel closed."); alertSubCh = nil }
				else { fmt.Printf("Alerts Subscriber 2: %s
", msg) }
			}
			if newsSubCh == nil && alertSubCh == nil {
				fmt.Println("News Subscriber 2: All channels closed, exiting.")
				return
			}
		}
	}()

	// Publisher sends messages
	time.Sleep(time.Millisecond * 100) // Give subscribers time to register
	broker.Publish("news", "Breaking News: Go channels are awesome!")
	broker.Publish("alerts", "Urgent: System maintenance tonight.")
	time.Sleep(time.Millisecond * 200)
	broker.Publish("news", "New feature release next week.")
	broker.Publish("alerts", "Reminder: Backups are running.")
	time.Sleep(time.Millisecond * 200)

	// Unsubscribe a subscriber (e.g., News Subscriber 1)
	// This requires keeping track of the channel returned by Subscribe for News Subscriber 1
	// For simplicity, let's assume we can't easily get it here directly without modifying the example.
	// In a real app, you'd store the returned 'newsSubCh' from `broker.Subscribe` in a variable.

	broker.Publish("news", "Final update: Server reboot complete.")
	broker.Publish("news", "The End") // Signal for a specific subscriber to stop

	// Give time for messages to be processed before main exits
	time.Sleep(time.Second * 1)
	// A more robust shutdown would involve signaling all subscribers to exit
	// and then closing the internal channels in the broker.
	wg.Wait() // Wait for subscribers to finish
	fmt.Println("Main: All messages published and subscribers processed.")

	// In a real system, you'd also need a way to close all subscriber channels 
	// managed by the broker when the broker itself shuts down.
}