package correlation

import (
	"context"
	"sync"
	"time"
)

// numShards est le nombre de partitions de la carte de sessions.
// Doit être une puissance de 2 pour permettre le masquage bitwise.
const numShards = 256

// shard est une partition thread-safe de la carte de sessions.
type shard struct {
	mu       sync.RWMutex
	sessions map[SessionKey]*SessionState
}

// Manager gère le cycle de vie des sessions TCP avec partitionnement
// pour réduire la contention lors des accès concurrents.
type Manager struct {
	shards  [numShards]shard
	timeout time.Duration // délai d'expiration des sessions (500ms par défaut)
	done    chan struct{}  // signal d'arrêt pour la goroutine GC

	// ReadyCh reçoit les sessions expirées prêtes à être écrites dans ClickHouse.
	ReadyCh chan *SessionState
}

// NewManager crée un Manager avec le délai d'expiration spécifié.
// Lance immédiatement les partitions de sessions.
func NewManager(timeout time.Duration) *Manager {
	m := &Manager{
		timeout: timeout,
		done:    make(chan struct{}),
		ReadyCh: make(chan *SessionState, 4096),
	}
	// Initialiser chaque shard
	for i := range m.shards {
		m.shards[i].sessions = make(map[SessionKey]*SessionState)
	}
	return m
}

// getShard retourne le shard correspondant à la clé de session.
// L'index est calculé par XOR des octets de l'IP et du port.
func (m *Manager) getShard(key SessionKey) *shard {
	idx := (key.SrcIP[3] ^ uint8(key.SrcPort>>8) ^ uint8(key.SrcPort)) & 0xFF
	return &m.shards[idx]
}

// GetOrCreate retourne la session existante ou en crée une nouvelle
// pour la clé donnée. Thread-safe.
func (m *Manager) GetOrCreate(key SessionKey) *SessionState {
	sh := m.getShard(key)

	// Essai rapide en lecture seule
	sh.mu.RLock()
	s, ok := sh.sessions[key]
	sh.mu.RUnlock()
	if ok {
		return s
	}

	// Création sous verrou exclusif
	sh.mu.Lock()
	defer sh.mu.Unlock()

	// Double-check après acquisition du verrou exclusif
	if s, ok = sh.sessions[key]; ok {
		return s
	}

	now := time.Now()
	s = &SessionState{
		Key:           key,
		FirstSeen:     now,
		LastActivity:  now,
		MaxKeepAlives: 1,
	}
	sh.sessions[key] = s
	return s
}

// Update applique la fonction fn sur la session identifiée par key,
// en créant la session si elle n'existe pas encore.
func (m *Manager) Update(key SessionKey, fn func(*SessionState)) {
	s := m.GetOrCreate(key)
	s.mu.Lock()
	defer s.mu.Unlock()
	fn(s)
	s.LastActivity = time.Now()
}

// StartGC lance la goroutine de nettoyage des sessions expirées.
// Toutes les 100ms, elle parcourt tous les shards et exporte
// les sessions expirées ou détectées comme Slowloris vers ReadyCh.
func (m *Manager) StartGC(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(100 * time.Millisecond)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-m.done:
				return
			case <-ticker.C:
				m.gcRound(10 * time.Second)
			}
		}
	}()
}

// gcRound effectue un passage de nettoyage sur tous les shards.
// Les sessions expirées ou Slowloris sont envoyées vers ReadyCh.
func (m *Manager) gcRound(slowlorisThreshold time.Duration) {
	for i := range m.shards {
		sh := &m.shards[i]

		// Collecter les clés à supprimer sans bloquer les écritures
		sh.mu.Lock()
		var toDelete []SessionKey
		for key, s := range sh.sessions {
			expired   := s.IsExpired(m.timeout)
			slowloris := s.IsSlowloris(slowlorisThreshold)

			if expired || slowloris {
				toDelete = append(toDelete, key)
				// Envoyer sans bloquer (drop si le canal est plein)
				select {
				case m.ReadyCh <- s:
				default:
				}
			}
		}
		for _, k := range toDelete {
			delete(sh.sessions, k)
		}
		sh.mu.Unlock()
	}
}

// Close arrête la goroutine GC et ferme le canal ReadyCh.
func (m *Manager) Close() {
	close(m.done)
	close(m.ReadyCh)
}
