package mcp

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"
)

const (
	partialSendIdleTTL       = time.Hour
	partialSendSweepInterval = time.Minute
	maxPartialSendEntries    = 1024
)

var errMessageServiceBusy = errors.New("message service is busy: too many active partial sends")

type partialSendKey struct {
	serverID     string
	localAgentID string
}

type partialSendEntry struct {
	gate           chan struct{}
	references     int
	lastUsed       time.Time
	pending        *pendingPartialSendState
	pendingTouched time.Time
}

type pendingPartialSendState struct {
	Target                string
	Signature             string
	LogicalClientID       string
	LastAcceptedSeq       int64
	NextChunkIndex        int
	AcceptedChunks        []acceptedSendChunk
	SettlementTargetKey   string
	SettlementSeenUpToSeq int64
	SettlementMessageID   string
}

// MessageService owns daemon-local MCP state shared by all handlers constructed
// by one Router. It is independent of wire.Protocol lifecycle and identity.
type MessageService struct {
	mu                  sync.Mutex
	partialSends        map[partialSendKey]*partialSendEntry
	partialSendCapacity int
	lastPartialSweep    time.Time
	now                 func() time.Time
	limits              *sendLimiterMap
}

// NewMessageService creates daemon-local state shared by all message handlers
// owned by one Router.
func NewMessageService() *MessageService {
	return newMessageServiceWithClock(time.Now)
}

func newMessageServiceWithClock(now func() time.Time) *MessageService {
	if now == nil {
		now = time.Now
	}
	return &MessageService{
		partialSends:        make(map[partialSendKey]*partialSendEntry),
		partialSendCapacity: maxPartialSendEntries,
		now:                 now,
		limits:              newsendLimiterMap(defaultSendRate, defaultSendBurst),
	}
}

type partialSendLease struct {
	service *MessageService
	key     partialSendKey
	entry   *partialSendEntry
	once    sync.Once
}

func (s *MessageService) acquirePartialSend(ctx context.Context, serverID string, localAgentID string) (*partialSendLease, error) {
	if s == nil {
		return nil, errors.New("message service is required")
	}
	key := partialSendKey{
		serverID:     strings.TrimSpace(serverID),
		localAgentID: strings.TrimSpace(localAgentID),
	}
	if key.serverID == "" || key.localAgentID == "" {
		return nil, errors.New("message service requires server and local agent identity")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	s.mu.Lock()
	now := s.now()
	s.sweepPartialSendsLocked(now)
	entry := s.partialSends[key]
	if entry == nil {
		if !s.makePartialSendRoomLocked(now) {
			s.mu.Unlock()
			return nil, errMessageServiceBusy
		}
		entry = &partialSendEntry{
			gate:     make(chan struct{}, 1),
			lastUsed: now,
		}
		entry.gate <- struct{}{}
		s.partialSends[key] = entry
	}
	entry.references++
	s.mu.Unlock()

	select {
	case <-ctx.Done():
		s.mu.Lock()
		entry.references--
		entry.lastUsed = s.now()
		s.mu.Unlock()
		return nil, ctx.Err()
	case <-entry.gate:
		return &partialSendLease{service: s, key: key, entry: entry}, nil
	}
}

func (s *MessageService) sweepPartialSendsLocked(now time.Time) {
	if !s.lastPartialSweep.IsZero() && now.Sub(s.lastPartialSweep) < partialSendSweepInterval {
		return
	}
	s.lastPartialSweep = now
	for key, entry := range s.partialSends {
		if entry.pending != nil && now.Sub(entry.pendingTouched) >= partialSendIdleTTL {
			entry.pending = nil
			entry.pendingTouched = time.Time{}
		}
		if entry.references == 0 && entry.pending == nil && now.Sub(entry.lastUsed) >= partialSendIdleTTL {
			delete(s.partialSends, key)
		}
	}
}

func (s *MessageService) makePartialSendRoomLocked(now time.Time) bool {
	if s.partialSendCapacity <= 0 || len(s.partialSends) < s.partialSendCapacity {
		return true
	}
	var (
		oldestKey   partialSendKey
		oldestEntry *partialSendEntry
	)
	for key, entry := range s.partialSends {
		if entry.pending != nil && now.Sub(entry.pendingTouched) >= partialSendIdleTTL {
			entry.pending = nil
			entry.pendingTouched = time.Time{}
		}
		if entry.references != 0 || entry.pending != nil {
			continue
		}
		if oldestEntry == nil || entry.lastUsed.Before(oldestEntry.lastUsed) ||
			(entry.lastUsed.Equal(oldestEntry.lastUsed) && partialSendKeyLess(key, oldestKey)) {
			oldestKey = key
			oldestEntry = entry
		}
	}
	if oldestEntry == nil {
		return false
	}
	delete(s.partialSends, oldestKey)
	return true
}

func partialSendKeyLess(left partialSendKey, right partialSendKey) bool {
	if left.serverID != right.serverID {
		return left.serverID < right.serverID
	}
	return left.localAgentID < right.localAgentID
}

func (l *partialSendLease) Pending() (pendingPartialSendState, bool) {
	if l == nil || l.service == nil || l.entry == nil {
		return pendingPartialSendState{}, false
	}
	l.service.mu.Lock()
	defer l.service.mu.Unlock()
	if l.entry.pending == nil {
		return pendingPartialSendState{}, false
	}
	now := l.service.now()
	if now.Sub(l.entry.pendingTouched) >= partialSendIdleTTL {
		l.entry.pending = nil
		l.entry.pendingTouched = time.Time{}
		return pendingPartialSendState{}, false
	}
	l.entry.pendingTouched = now
	l.entry.lastUsed = now
	pending := *l.entry.pending
	pending.AcceptedChunks = append([]acceptedSendChunk(nil), pending.AcceptedChunks...)
	return pending, true
}

func (l *partialSendLease) StorePending(pending pendingPartialSendState) {
	if l == nil || l.service == nil || l.entry == nil {
		return
	}
	l.service.mu.Lock()
	defer l.service.mu.Unlock()
	now := l.service.now()
	stored := pending
	stored.AcceptedChunks = append([]acceptedSendChunk(nil), pending.AcceptedChunks...)
	l.entry.pending = &stored
	l.entry.pendingTouched = now
	l.entry.lastUsed = now
}

func (l *partialSendLease) ClearPending() {
	if l == nil || l.service == nil || l.entry == nil {
		return
	}
	l.service.mu.Lock()
	defer l.service.mu.Unlock()
	l.entry.pending = nil
	l.entry.pendingTouched = time.Time{}
	l.entry.lastUsed = l.service.now()
}

func (l *partialSendLease) Release() {
	if l == nil {
		return
	}
	l.once.Do(func() {
		l.service.mu.Lock()
		l.entry.references--
		l.entry.lastUsed = l.service.now()
		l.service.mu.Unlock()
		l.entry.gate <- struct{}{}
	})
}
