package packets

import (
	"sync"
	"time"
)

type SessionManager struct {
	sessions map[uint32]*Session
	mu       sync.RWMutex
}

func NewSessionManager() *SessionManager {
	return &SessionManager{
		sessions: make(map[uint32]*Session),
	}
}

func (sm *SessionManager) GetSession(id uint32) *Session {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return sm.sessions[id]
}

func (sm *SessionManager) CreateSession(id uint32) *Session {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	session := &Session{
		ID:          id,
		Established: false,
		NextSeq:     0,
		LastKeep:    time.Now(),
		Compressed:  false,
		Fragments:   make(map[uint16]*FragmentBuffer),
	}

	sm.sessions[id] = session
	return session
}

func (sm *SessionManager) RemoveSession(id uint32) {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	delete(sm.sessions, id)
}

func (sm *SessionManager) UpdateKeepAlive(id uint32) {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	if session := sm.sessions[id]; session != nil {
		session.LastKeep = time.Now()
	}
}

func (sm *SessionManager) GetNextSequence(id uint32) uint16 {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if session := sm.sessions[id]; session != nil {
		seq := session.NextSeq
		session.NextSeq++
		return seq
	}
	return 0
}

func (sm *SessionManager) SetEstablished(id uint32, established bool) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if session := sm.sessions[id]; session != nil {
		session.Established = established
	}
}