package webui

import (
	"net/http"
	"sync"

	"github.com/ProtonMail/go-crypto/openpgp"
)

// pgpKeyCache holds unlocked PGP identities for the rest of a login session —
// mirrors smimeKeyCache (webmail_smime_cache.go), but caches the whole *openpgp.Entity
// (not just a raw private key) since decrypting a PGP message needs subkey lookup on
// the entity itself, not a bare key value. Kept as its own type rather than
// generalizing smimeKeyCache into an any-typed cache — both stay simply and
// correctly typed.
type pgpKeyCache struct {
	mu    sync.Mutex
	byTok map[string]map[int64]*openpgp.Entity
}

func newPGPKeyCache() *pgpKeyCache {
	return &pgpKeyCache{byTok: map[string]map[int64]*openpgp.Entity{}}
}

func (c *pgpKeyCache) get(token string, identityID int64) (*openpgp.Entity, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	keys, ok := c.byTok[token]
	if !ok {
		return nil, false
	}
	entity, ok := keys[identityID]
	return entity, ok
}

func (c *pgpKeyCache) put(token string, identityID int64, entity *openpgp.Entity) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.byTok[token] == nil {
		c.byTok[token] = map[int64]*openpgp.Entity{}
	}
	c.byTok[token][identityID] = entity
}

// clearSession drops every unlocked identity for one session — called on logout so
// a key never outlives the session it was unlocked in. Same ponytail-flagged
// simplification as smimeKeyCache.clearSession: memory-bounded by active sessions,
// not TTL'd against a session that expires without an explicit logout.
func (c *pgpKeyCache) clearSession(token string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.byTok, token)
}

// sessionToken reads the raw webmail session cookie value, used as the PGP key
// cache's key — separate from mailboxFromContext, which only exposes the resolved
// *db.Mailbox, not the token itself.
func sessionToken(r *http.Request) string {
	c, err := r.Cookie(mailboxSessionCookieName)
	if err != nil {
		return ""
	}
	return c.Value
}