package middleware

import (
	"maps"
	"sync"
	"sync/atomic"
	"time"
)

// DedupMode controls what content is sent for previously-enriched tables.
type DedupMode string

const (
	// DedupModeReference sends a minimal reference (table name, description, note).
	DedupModeReference DedupMode = "reference"

	// DedupModeSummary sends table-level semantic_context without column_context.
	DedupModeSummary DedupMode = "summary"

	// DedupModeNone sends no enrichment at all for repeat queries.
	DedupModeNone DedupMode = "none"
)

// SentTableEntry records when a table was enriched and the estimated token count.
type SentTableEntry struct {
	SentAt     time.Time `json:"sent_at"`
	TokenCount int       `json:"token_count"`
}

// SessionEnrichmentCache tracks which tables have been enriched per client
// session, enabling deduplication of large semantic metadata blocks.
type SessionEnrichmentCache struct {
	mu             sync.RWMutex
	sessions       map[string]*sessionState
	entryTTL       time.Duration
	sessionTimeout time.Duration
	done           chan struct{}

	// Cumulative token counters for diagnostics.
	tokensFull    atomic.Int64
	tokensDeduped atomic.Int64
}

// sessionState tracks enrichment state for a single session.
type sessionState struct {
	sentTables map[string]SentTableEntry
	lastAccess time.Time
}

// NewSessionEnrichmentCache creates a new session enrichment cache.
func NewSessionEnrichmentCache(entryTTL, sessionTimeout time.Duration) *SessionEnrichmentCache {
	return &SessionEnrichmentCache{
		sessions:       make(map[string]*sessionState),
		entryTTL:       entryTTL,
		sessionTimeout: sessionTimeout,
		done:           make(chan struct{}),
	}
}

// MarkSent records that full enrichment was sent for the given table in this session.
func (c *SessionEnrichmentCache) MarkSent(sessionID, tableKey string, tokenCount int) {
	c.mu.Lock()
	defer c.mu.Unlock()

	state := c.getOrCreateSession(sessionID)
	state.sentTables[tableKey] = SentTableEntry{
		SentAt:     time.Now(),
		TokenCount: tokenCount,
	}
	state.lastAccess = time.Now()
}

// WasSentRecently returns true if the table was enriched within the entry TTL
// for this session.
func (c *SessionEnrichmentCache) WasSentRecently(sessionID, tableKey string) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()

	state, ok := c.sessions[sessionID]
	if !ok {
		return false
	}

	entry, ok := state.sentTables[tableKey]
	if !ok {
		return false
	}

	return time.Since(entry.SentAt) < c.entryTTL
}

// GetTokenCount returns the stored full enrichment token count for a table
// in a session, or 0 if not found.
func (c *SessionEnrichmentCache) GetTokenCount(sessionID, tableKey string) int {
	c.mu.RLock()
	defer c.mu.RUnlock()

	state, ok := c.sessions[sessionID]
	if !ok {
		return 0
	}

	entry, ok := state.sentTables[tableKey]
	if !ok {
		return 0
	}

	return entry.TokenCount
}

// AddTokensFull adds to the cumulative full enrichment token counter.
func (c *SessionEnrichmentCache) AddTokensFull(n int64) {
	c.tokensFull.Add(n)
}

// AddTokensDeduped adds to the cumulative deduped enrichment token counter.
func (c *SessionEnrichmentCache) AddTokensDeduped(n int64) {
	c.tokensDeduped.Add(n)
}

// TokensFull returns the cumulative full enrichment tokens sent.
func (c *SessionEnrichmentCache) TokensFull() int64 {
	return c.tokensFull.Load()
}

// TokensDeduped returns the cumulative deduped enrichment tokens sent.
func (c *SessionEnrichmentCache) TokensDeduped() int64 {
	return c.tokensDeduped.Load()
}

// StartCleanup starts a background goroutine that evicts idle sessions.
func (c *SessionEnrichmentCache) StartCleanup(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()

		for {
			select {
			case <-c.done:
				return
			case <-ticker.C:
				c.cleanup()
			}
		}
	}()
}

// Stop stops the background cleanup goroutine.
func (c *SessionEnrichmentCache) Stop() {
	close(c.done)
}

// SessionCount returns the number of tracked sessions.
func (c *SessionEnrichmentCache) SessionCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.sessions)
}

// EntryTTL returns the configured entry TTL for diagnostics.
func (c *SessionEnrichmentCache) EntryTTL() time.Duration {
	return c.entryTTL
}

// getOrCreateSession returns the session state, creating it if needed.
// Caller must hold the write lock.
func (c *SessionEnrichmentCache) getOrCreateSession(sessionID string) *sessionState {
	state, ok := c.sessions[sessionID]
	if !ok {
		state = &sessionState{
			sentTables: make(map[string]SentTableEntry),
			lastAccess: time.Now(),
		}
		c.sessions[sessionID] = state
	}
	return state
}

// LoadSession pre-populates dedup state for a session from external storage.
// This is used on startup to restore enrichment dedup continuity from the
// session store.
func (c *SessionEnrichmentCache) LoadSession(sessionID string, sentTables map[string]SentTableEntry) {
	c.mu.Lock()
	defer c.mu.Unlock()

	state := c.getOrCreateSession(sessionID)
	maps.Copy(state.sentTables, sentTables)
	state.lastAccess = time.Now()
}

// ExportSessions returns all session dedup states for persistence.
// Each map entry maps a session ID to its sent-table entries.
func (c *SessionEnrichmentCache) ExportSessions() map[string]map[string]SentTableEntry {
	c.mu.RLock()
	defer c.mu.RUnlock()

	result := make(map[string]map[string]SentTableEntry, len(c.sessions))
	for id, state := range c.sessions {
		tables := make(map[string]SentTableEntry, len(state.sentTables))
		maps.Copy(tables, state.sentTables)
		result[id] = tables
	}
	return result
}

// cleanup evicts sessions that have been idle longer than sessionTimeout,
// and removes expired entries from active sessions.
func (c *SessionEnrichmentCache) cleanup() {
	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	for id, state := range c.sessions {
		if now.Sub(state.lastAccess) > c.sessionTimeout {
			delete(c.sessions, id)
			continue
		}
		// Remove expired entries from active sessions
		for table, entry := range state.sentTables {
			if now.Sub(entry.SentAt) > c.entryTTL {
				delete(state.sentTables, table)
			}
		}
	}
}
