// Package session manages secure communication sessions with Forward Secrecy
package session

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"sync"
	"time"

	"github.com/flynn/noise"
)

// Session represents a secure communication session
type Session struct {
	ID           string
	RemotePubKey string // Remote peer's public key (hex)
	HandshakeState *noise.HandshakeState
	CipherState  *noise.CipherState
	CreatedAt    time.Time
	LastActivity time.Time
	IsInitiator  bool
}

// SessionManager manages multiple sessions
type SessionManager struct {
	sessions map[string]*Session // key = session ID
	mu       sync.RWMutex
}

// NewSessionManager creates a new session manager
func NewSessionManager() *SessionManager {
	return &SessionManager{
		sessions: make(map[string]*Session),
	}
}

// CreateSession creates a new Noise Protocol session
func (sm *SessionManager) CreateSession(remotePubKey string, isInitiator bool) (*Session, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	// Generate session ID
	sessionID := generateSessionID()

	// Create Noise handshake configuration
	// Using Noise_XX_25519_ChaChaPoly_BLAKE2b pattern for mutual authentication
	var hs *noise.HandshakeState
	var err error

	config := noise.Config{
		CipherSuite: noise.NewCipherSuite(noise.DH25519, noise.CipherChaChaPoly, noise.HashBLAKE2b),
		Pattern:     noise.HandshakeXX,
		Initiator:   isInitiator,
	}

	// Generate ephemeral keypair
	keypair, err := noise.DH25519.GenerateKeypair(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate keypair: %w", err)
	}

	config.StaticKeypair = keypair

	hs, err = noise.NewHandshakeState(config)
	if err != nil {
		return nil, fmt.Errorf("failed to create handshake state: %w", err)
	}

	session := &Session{
		ID:             sessionID,
		RemotePubKey:   remotePubKey,
		HandshakeState: hs,
		CreatedAt:      time.Now(),
		LastActivity:   time.Now(),
		IsInitiator:    isInitiator,
	}

	sm.sessions[sessionID] = session
	return session, nil
}

// GetSession retrieves a session by ID
func (sm *SessionManager) GetSession(sessionID string) (*Session, bool) {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	session, exists := sm.sessions[sessionID]
	return session, exists
}

// RemoveSession removes a session and securely wipes it
func (sm *SessionManager) RemoveSession(sessionID string) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if session, exists := sm.sessions[sessionID]; exists {
		session.SecureWipe()
		delete(sm.sessions, sessionID)
	}
}

// SecureWipeAll securely wipes all sessions
func (sm *SessionManager) SecureWipeAll() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	for _, session := range sm.sessions {
		session.SecureWipe()
	}
	sm.sessions = make(map[string]*Session)
}

// ListSessions returns all active sessions
func (sm *SessionManager) ListSessions() []*Session {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	sessions := make([]*Session, 0, len(sm.sessions))
	for _, session := range sm.sessions {
		sessions = append(sessions, session)
	}
	return sessions
}

// CleanupExpiredSessions removes sessions older than the specified duration
func (sm *SessionManager) CleanupExpiredSessions(maxAge time.Duration) int {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	now := time.Now()
	removed := 0

	for id, session := range sm.sessions {
		if now.Sub(session.LastActivity) > maxAge {
			session.SecureWipe()
			delete(sm.sessions, id)
			removed++
		}
	}

	return removed
}

// UpdateActivity updates the last activity timestamp
func (sm *SessionManager) UpdateActivity(sessionID string) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	if session, exists := sm.sessions[sessionID]; exists {
		session.LastActivity = time.Now()
	}
}

// Count returns the number of active sessions
func (sm *SessionManager) Count() int {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return len(sm.sessions)
}

// SecureWipe zeros out sensitive session data
func (s *Session) SecureWipe() {
	// Wipe cipher state if exists
	if s.CipherState != nil {
		// Noise library doesn't expose key wiping, but we can nil the reference
		s.CipherState = nil
	}

	// Wipe handshake state
	if s.HandshakeState != nil {
		s.HandshakeState = nil
	}

	// Zero out session ID
	s.ID = ""
	s.RemotePubKey = ""
}

// generateSessionID generates a random session ID
func generateSessionID() string {
	id := make([]byte, 16)
	rand.Read(id)
	return hex.EncodeToString(id)
}

// String returns a safe string representation
func (s *Session) String() string {
	role := "responder"
	if s.IsInitiator {
		role = "initiator"
	}
	return fmt.Sprintf("Session(id=%s, role=%s, remote=%s...)",
		s.ID[:8], role, s.RemotePubKey[:16])
}

// String returns session manager info
func (sm *SessionManager) String() string {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return fmt.Sprintf("SessionManager(active=%d)", len(sm.sessions))
}