package mcp

import (
	"strings"
	"sync"
	"time"
)

// Default outbound rate limit: 5 messages/second with a burst of 8.
// The burst accommodates chunked long messages (split into 4000-char chunks)
// plus attachments in a single tool call.
const (
	defaultSendRate  = 5.0 // tokens per second
	defaultSendBurst = 8   // maximum burst size

	sendLimiterIdleTTL         = time.Hour
	sendLimiterCleanupInterval = time.Minute
	maxsendLimiterBuckets      = 1024
)

// sendLimiter enforces per-session outbound message rate limits using a
// token-bucket algorithm. All methods are safe for concurrent use.
type sendLimiter struct {
	mu       sync.Mutex
	tokens   float64
	max      float64   // burst capacity
	rate     float64   // tokens per second (refill rate)
	lastTime time.Time // last refill time
}

// newsendLimiter creates a sendLimiter allowing rate messages/sec with a burst
// of burst. The bucket starts full (burst tokens available).
func newsendLimiter(rate float64, burst int) *sendLimiter {
	return &sendLimiter{
		tokens:   float64(burst),
		max:      float64(burst),
		rate:     rate,
		lastTime: time.Now(),
	}
}

// Allow reports whether an outbound send is permitted and, if so, consumes
// one token. Returns false when the rate limit has been exceeded.
func (l *sendLimiter) Allow() bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(l.lastTime).Seconds()
	l.tokens += elapsed * l.rate
	if l.tokens > l.max {
		l.tokens = l.max
	}
	l.lastTime = now

	if l.tokens < 1 {
		return false
	}
	l.tokens--
	return true
}

// sendLimiterMap manages per-session sendLimiters. All methods are safe for concurrent use.
type sendLimiterMap struct {
	mu          sync.Mutex
	limiters    map[string]sendLimiterEntry
	rate        float64
	burst       int
	maxBuckets  int
	lastCleanup time.Time
}

type sendLimiterEntry struct {
	limiter  *sendLimiter
	lastUsed time.Time
}

// newsendLimiterMap creates a sendLimiterMap with the given rate and burst.
func newsendLimiterMap(rate float64, burst int) *sendLimiterMap {
	return &sendLimiterMap{
		limiters:   make(map[string]sendLimiterEntry),
		rate:       rate,
		burst:      burst,
		maxBuckets: maxsendLimiterBuckets,
	}
}

// Get returns the sendLimiter for the given session key, creating one if needed.
func (m *sendLimiterMap) Get(key string) *sendLimiter {
	key = strings.TrimSpace(key)
	if key == "" {
		key = "default"
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	now := time.Now()
	m.removeStaleLocked(now)
	entry, ok := m.limiters[key]
	if !ok {
		m.removeOldestUntilRoomLocked()
		entry = sendLimiterEntry{limiter: newsendLimiter(m.rate, m.burst)}
	}
	entry.lastUsed = now
	m.limiters[key] = entry
	return entry.limiter
}

// Remove deletes the limiter for the given key (e.g. on session end).
func (m *sendLimiterMap) Remove(key string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.limiters, key)
}

func (m *sendLimiterMap) removeStaleLocked(now time.Time) {
	if !m.lastCleanup.IsZero() && now.Sub(m.lastCleanup) < sendLimiterCleanupInterval {
		return
	}
	m.lastCleanup = now
	for key, entry := range m.limiters {
		if now.Sub(entry.lastUsed) > sendLimiterIdleTTL {
			delete(m.limiters, key)
		}
	}
}

func (m *sendLimiterMap) removeOldestUntilRoomLocked() {
	if m.maxBuckets <= 0 {
		return
	}
	for len(m.limiters) >= m.maxBuckets {
		var oldestKey string
		var oldestTime time.Time
		for key, entry := range m.limiters {
			if oldestKey == "" || entry.lastUsed.Before(oldestTime) {
				oldestKey = key
				oldestTime = entry.lastUsed
			}
		}
		if oldestKey == "" {
			return
		}
		delete(m.limiters, oldestKey)
	}
}

// sendLimiterKey returns a stable per-conversation key for outbound send
// limiting. If serverID or conversationKey is unavailable, it falls back to
// agentID so call sites with only agent-scoped state remain bounded.
func sendLimiterKey(serverID, conversationKey, agentID string) string {
	serverID = strings.TrimSpace(serverID)
	conversationKey = strings.TrimSpace(conversationKey)
	if serverID != "" && conversationKey != "" {
		return serverID + "\x00" + conversationKey
	}
	if conversationKey != "" {
		return conversationKey
	}
	return strings.TrimSpace(agentID)
}