package security

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"time"
)

type SessionManager struct {
	cookieName string
	secure     bool
	sameSite   http.SameSite
	maxAge     time.Duration
	aead       cipher.AEAD
}

func NewSessionManager(secret []byte, cookieName string, secure bool) (*SessionManager, error) {
	// Derive 32-byte key for AES-256-GCM
	key := hmacSHA256(secret, []byte("ntfywui session v1"))
	if len(key) != 32 {
		return nil, io.ErrUnexpectedEOF
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &SessionManager{
		cookieName: cookieName,
		secure:     secure,
		sameSite:   http.SameSiteLaxMode,
		maxAge:     12 * time.Hour,
		aead:       aead,
	}, nil
}

// Session contents are encrypted+authenticated.
type Session struct {
	User      string `json:"user"`
	Role      string `json:"role"`
	CSRF      string `json:"csrf"`
	Flash     string `json:"flash,omitempty"`
	IssuedAt  int64  `json:"iat"`
	ExpiresAt int64  `json:"exp"`
}

func (sm *SessionManager) Get(r *http.Request) (*Session, bool) {
	c, err := r.Cookie(sm.cookieName)
	if err != nil || c.Value == "" {
		return &Session{}, false
	}
	raw, err := base64.RawURLEncoding.DecodeString(c.Value)
	if err != nil || len(raw) < sm.aead.NonceSize() {
		return &Session{}, false
	}
	nonce := raw[:sm.aead.NonceSize()]
	ct := raw[sm.aead.NonceSize():]
	pt, err := sm.aead.Open(nil, nonce, ct, nil)
	if err != nil {
		return &Session{}, false
	}
	var s Session
	if err := json.Unmarshal(pt, &s); err != nil {
		return &Session{}, false
	}
	now := time.Now().Unix()
	if s.ExpiresAt != 0 && now > s.ExpiresAt {
		return &Session{}, false
	}
	return &s, true
}

func (sm *SessionManager) Save(w http.ResponseWriter, s *Session) error {
	now := time.Now()
	if s.IssuedAt == 0 {
		s.IssuedAt = now.Unix()
	}
	if s.ExpiresAt == 0 {
		s.ExpiresAt = now.Add(sm.maxAge).Unix()
	}
	pt, err := json.Marshal(s)
	if err != nil {
		return err
	}
	nonce := make([]byte, sm.aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return err
	}
	ct := sm.aead.Seal(nil, nonce, pt, nil)
	raw := append(nonce, ct...)
	val := base64.RawURLEncoding.EncodeToString(raw)
	http.SetCookie(w, &http.Cookie{
		Name:     sm.cookieName,
		Value:    val,
		Path:     "/",
		HttpOnly: true,
		Secure:   sm.secure,
		SameSite: sm.sameSite,
		MaxAge:   int(sm.maxAge.Seconds()),
	})
	return nil
}

func (sm *SessionManager) Clear(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     sm.cookieName,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   sm.secure,
		SameSite: sm.sameSite,
		MaxAge:   -1,
	})
}