package main

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

// Message is the data type for messages in our Pub/Sub system.
type Message struct {
	Topic   string
	Payload string
}

// PubSub represents our simplified publish-subscribe system.
type PubSub struct {
	mu          sync.RWMutex
	subscribers map[string][]chan Message // topic -> list of subscriber channels
}

// NewPubSub creates and returns a new PubSub instance.
func NewPubSub() *PubSub {
	return &PubSub{
		subscribers: make(map[string][]chan Message),
	}
}

// Subscribe allows a client to subscribe to a given topic.
// It returns a read-only channel where messages for that topic will be delivered.
func (ps *PubSub) Subscribe(topic string) <-chan Message {
	ps.mu.Lock()
	defer ps.mu.Unlock()

	ch := make(chan Message, 10) // Buffered channel for subscriber
	ps.subscribers[topic] = append(ps.subscribers[topic], ch)
	fmt.Printf("Subscribed to topic '%s'. Total subscribers: %d
", topic, len(ps.subscribers[topic]))
	return ch
}

// Unsubscribe removes a subscriber channel from a topic.
func (ps *PubSub) Unsubscribe(topic string, subCh <-chan Message) {
	ps.mu.Lock()
	defer ps.mu.Unlock()

	if channels, ok := ps.subscribers[topic]; ok {
		for i, ch := range channels {
			if ch == subCh {
				// Remove the channel from the slice
				ps.subscribers[topic] = append(channels[:i], channels[i+1:]...)
				close(ch) // Close the subscriber's channel
				fmt.Printf("Unsubscribed from topic '%s'. Remaining subscribers: %d
", topic, len(ps.subscribers[topic]))
				return
			}
		}
	}
}

// Publish sends a message to all subscribers of a specific topic.
func (ps *PubSub) Publish(msg Message) {
	ps.mu.RLock() // Use RLock because we are only reading the map for subscribers
	defer ps.mu.RUnlock()

	if channels, ok := ps.subscribers[msg.Topic]; ok {
		// Send message to all subscribers in separate goroutines to avoid blocking
		// if a subscriber channel is full or slow.
		for _, ch := range channels {
			go func(c chan Message) {
				select {
				case c <- msg:
					// Message sent
				case <-time.After(100 * time.Millisecond):
					// Timeout if subscriber channel is full.
					// This prevents a slow subscriber from blocking others indefinitely.
					fmt.Printf("Warning: Dropped message for topic '%s' due to slow subscriber.
", msg.Topic)
				}
			}(ch)
		}
	} else {
		fmt.Printf("No subscribers for topic '%s'. Message dropped.
", msg.Topic)
	}
}

func main() {
	ps := NewPubSub()

	// Subscriber 1
	sub1 := ps.Subscribe("news")
	go func() {
		for msg := range sub1 {
			fmt.Printf("Subscriber 1 received: [%s] %s
", msg.Topic, msg.Payload)
		}
		fmt.Println("Subscriber 1 channel closed.")
	}()

	// Subscriber 2
	sub2 := ps.Subscribe("news")
	go func() {
		for msg := range sub2 {
			fmt.Printf("Subscriber 2 received: [%s] %s
", msg.Topic, msg.Payload)
			time.Sleep(200 * time.Millisecond) // Simulate slow subscriber
		}
		fmt.Println("Subscriber 2 channel closed.")
	}()

	// Subscriber 3 for a different topic
	sub3 := ps.Subscribe("updates")
	go func() {
		for msg := range sub3 {
			fmt.Printf("Subscriber 3 received: [%s] %s
", msg.Topic, msg.Payload)
		}
		fmt.Println("Subscriber 3 channel closed.")
	}()

	time.Sleep(100 * time.Millisecond) // Give subscribers time to start

	ps.Publish(Message{Topic: "news", Payload: "Breaking news: Gophers spotted!"})
	ps.Publish(Message{Topic: "updates", Payload: "System update available."})
	ps.Publish(Message{Topic: "news", Payload: "Local weather: Sunny with a chance of channels."})
	ps.Publish(Message{Topic: "alerts", Payload: "Emergency alert!"}) // No subscriber for 'alerts'

	time.Sleep(1 * time.Second) // Let messages be processed

	ps.Unsubscribe("news", sub1) // Subscriber 1 leaves
	ps.Publish(Message{Topic: "news", Payload: "More news after unsubscribe."})

	time.Sleep(500 * time.Millisecond)
	fmt.Println("Main: Exiting.")
	// In a real application, you'd have a mechanism to close all subscriber channels
	// when the PubSub system shuts down. For this example, we let main exit.
}