package memory

import (
	"context"
	"sync"

	"github.com/rwadurian/mpc-system/pkg/logger"
	"github.com/rwadurian/mpc-system/services/message-router/application/use_cases"
	"github.com/rwadurian/mpc-system/services/message-router/domain/entities"
	"go.uber.org/zap"
)

// MessageBrokerAdapter implements MessageBroker using in-memory channels
type MessageBrokerAdapter struct {
	mu sync.RWMutex
	// partyChannels maps partyID to their message channel
	partyChannels map[string]chan *entities.MessageDTO
	// sessionChannels maps sessionID:partyID to their broadcast channel
	sessionChannels map[string]chan *entities.MessageDTO
}

// NewMessageBrokerAdapter creates a new in-memory message broker
func NewMessageBrokerAdapter() *MessageBrokerAdapter {
	return &MessageBrokerAdapter{
		partyChannels:   make(map[string]chan *entities.MessageDTO),
		sessionChannels: make(map[string]chan *entities.MessageDTO),
	}
}

// PublishToParty publishes a message to a specific party
func (a *MessageBrokerAdapter) PublishToParty(ctx context.Context, partyID string, message *entities.MessageDTO) error {
	a.mu.RLock()
	ch, exists := a.partyChannels[partyID]
	a.mu.RUnlock()

	if !exists {
		logger.Debug("party channel not found, message will be available via polling",
			zap.String("party_id", partyID))
		return nil
	}

	select {
	case ch <- message:
		logger.Debug("published message to party",
			zap.String("party_id", partyID),
			zap.String("message_id", message.ID))
	case <-ctx.Done():
		return ctx.Err()
	default:
		// Channel full, message will be available via polling
		logger.Warn("party channel full, message available via polling",
			zap.String("party_id", partyID))
	}

	return nil
}

// PublishToSession publishes a message to all parties in a session (except sender)
func (a *MessageBrokerAdapter) PublishToSession(
	ctx context.Context,
	sessionID string,
	excludeParty string,
	message *entities.MessageDTO,
) error {
	a.mu.RLock()
	defer a.mu.RUnlock()

	prefix := sessionID + ":"

	// Debug: log all session channels
	var allKeys []string
	for key := range a.sessionChannels {
		allKeys = append(allKeys, key)
	}
	logger.Debug("PublishToSession looking for channels",
		zap.String("session_id", sessionID),
		zap.String("exclude_party", excludeParty),
		zap.String("prefix", prefix),
		zap.Int("total_session_channels", len(a.sessionChannels)),
		zap.Strings("all_keys", allKeys))

	foundCount := 0
	for key, ch := range a.sessionChannels {
		if len(key) > len(prefix) && key[:len(prefix)] == prefix {
			partyID := key[len(prefix):]
			if partyID == excludeParty {
				logger.Debug("skipping sender party",
					zap.String("session_id", sessionID),
					zap.String("party_id", partyID))
				continue
			}

			foundCount++
			select {
			case ch <- message:
				logger.Debug("broadcast message to party",
					zap.String("session_id", sessionID),
					zap.String("party_id", partyID))
			case <-ctx.Done():
				return ctx.Err()
			default:
				// Channel full
				logger.Warn("session channel full for party",
					zap.String("party_id", partyID))
			}
		}
	}

	if foundCount == 0 {
		logger.Warn("No parties found for broadcast",
			zap.String("session_id", sessionID),
			zap.String("exclude_party", excludeParty))
	}

	return nil
}

// SubscribeToPartyMessages subscribes to messages for a specific party
// If the party already has an active subscription, the old channel is closed first
func (a *MessageBrokerAdapter) SubscribeToPartyMessages(
	ctx context.Context,
	partyID string,
) (<-chan *entities.MessageDTO, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	// Close existing channel if party is re-subscribing (e.g., after reconnect)
	if oldCh, exists := a.partyChannels[partyID]; exists {
		close(oldCh)
		logger.Info("closed existing party channel for re-subscription",
			zap.String("party_id", partyID))
	}

	// Create new channel
	a.partyChannels[partyID] = make(chan *entities.MessageDTO, 100)
	ch := a.partyChannels[partyID]

	// Return a read-only channel
	out := make(chan *entities.MessageDTO, 100)
	go func() {
		defer close(out)
		for {
			select {
			case <-ctx.Done():
				return
			case msg, ok := <-ch:
				if !ok {
					return
				}
				select {
				case out <- msg:
				case <-ctx.Done():
					return
				}
			}
		}
	}()

	return out, nil
}

// SubscribeToSessionMessages subscribes to all messages in a session
// If the party already has an active subscription for this session, the old channel is closed first
func (a *MessageBrokerAdapter) SubscribeToSessionMessages(
	ctx context.Context,
	sessionID string,
	partyID string,
) (<-chan *entities.MessageDTO, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	key := sessionID + ":" + partyID

	logger.Info("SubscribeToSessionMessages called",
		zap.String("session_id", sessionID),
		zap.String("party_id", partyID),
		zap.String("key", key),
		zap.Int("current_channel_count", len(a.sessionChannels)))

	// Close existing channel if party is re-subscribing (e.g., after reconnect)
	if oldCh, exists := a.sessionChannels[key]; exists {
		close(oldCh)
		logger.Info("closed existing session channel for re-subscription",
			zap.String("key", key))
	}

	// Create new channel
	a.sessionChannels[key] = make(chan *entities.MessageDTO, 100)
	ch := a.sessionChannels[key]
	logger.Info("Created new session channel",
		zap.String("key", key))

	// Return a read-only channel
	out := make(chan *entities.MessageDTO, 100)
	go func() {
		defer close(out)
		for {
			select {
			case <-ctx.Done():
				// Cleanup on context done
				a.mu.Lock()
				delete(a.sessionChannels, key)
				a.mu.Unlock()
				return
			case msg, ok := <-ch:
				if !ok {
					return
				}
				select {
				case out <- msg:
				case <-ctx.Done():
					return
				}
			}
		}
	}()

	return out, nil
}

// UnsubscribeParty removes a party's channel
func (a *MessageBrokerAdapter) UnsubscribeParty(partyID string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	if ch, exists := a.partyChannels[partyID]; exists {
		close(ch)
		delete(a.partyChannels, partyID)
	}
}

// Close closes all channels
func (a *MessageBrokerAdapter) Close() error {
	a.mu.Lock()
	defer a.mu.Unlock()

	for _, ch := range a.partyChannels {
		close(ch)
	}
	for _, ch := range a.sessionChannels {
		close(ch)
	}

	a.partyChannels = make(map[string]chan *entities.MessageDTO)
	a.sessionChannels = make(map[string]chan *entities.MessageDTO)

	return nil
}

// Ensure interface compliance
var _ use_cases.MessageBroker = (*MessageBrokerAdapter)(nil)
