// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0

// Package session implements HMAC-signed cookie sessions for the admin
// UI. The cookie is the only state: it is signed, verified and expired
// here, and no session table exists anywhere.
//
// The cookie value is base64url(payload) + "." + base64url(signature)
// where payload is JSON {"d": {…}, "iat": unix-seconds}. Sessions are
// invalidated by signature mismatch or expiry, so users simply log in
// again after the signing key changes.
package session

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

// CookieName is the session cookie name.
const CookieName = "volumen_session"

type contextKey struct{}

// Store signs and verifies session cookies.
type Store struct {
	secret []byte
	ttl    time.Duration
	secure bool
}

// New creates a session store. An empty secret is replaced with an
// ephemeral random one (sessions then do not survive restarts).
func New(secret string, ttl time.Duration, secure bool) *Store {
	key := []byte(secret)
	if len(key) == 0 {
		key = make([]byte, 64)
		if _, err := rand.Read(key); err != nil {
			slog.Warn("session: random source failed", "error", err)
		}
		slog.Warn("session: session_key is empty; using an ephemeral random secret")
	}
	return &Store{secret: key, ttl: ttl, secure: secure}
}

// Secure reports whether cookies carry the Secure flag.
func (s *Store) Secure() bool { return s.secure }

// Session is a mutable key-value bag loaded for one request.
type Session struct {
	data  map[string]string
	dirty bool
	store *Store
}

// Get returns the value for key, or "".
func (s *Session) Get(key string) string { return s.data[key] }

// Set stores a value and marks the session for re-signing.
func (s *Session) Set(key, value string) {
	if s.data[key] == value {
		return
	}
	s.data[key] = value
	s.dirty = true
}

// Delete removes a key.
func (s *Session) Delete(key string) {
	if _, ok := s.data[key]; !ok {
		return
	}
	delete(s.data, key)
	s.dirty = true
}

// Clear empties the session.
func (s *Session) Clear() {
	if len(s.data) == 0 {
		return
	}
	s.data = map[string]string{}
	s.dirty = true
}

// Abandon empties the session and clears the dirty flag, for a response
// that expires the cookie instead of re-signing it. Save would otherwise
// append a fresh cookie after Destroy's expiring one, and the browser
// applies the last header.
func (s *Session) Abandon() {
	s.data = map[string]string{}
	s.dirty = false
}

type payload struct {
	Data map[string]string `json:"d"`
	Iat  int64             `json:"iat"`
}

// Load reads and verifies the session cookie from the request; a
// missing or invalid cookie yields an empty session.
func (s *Store) Load(r *http.Request) *Session {
	sess := &Session{data: map[string]string{}, store: s}
	cookie, err := r.Cookie(CookieName)
	if err != nil || cookie.Value == "" {
		return sess
	}
	data, ok := s.verify(cookie.Value)
	if !ok {
		return sess
	}
	var p payload
	if err := json.Unmarshal(data, &p); err != nil {
		return sess
	}
	if s.ttl > 0 && time.Since(time.Unix(p.Iat, 0)) > s.ttl {
		return sess
	}
	if p.Data != nil {
		sess.data = p.Data
	} else {
		sess.data = map[string]string{}
	}
	return sess
}

func (s *Store) verify(value string) ([]byte, bool) {
	// The signature is appended after the last dot; the body is base64url
	// and carries none.
	body, sig, found := strings.CutLast(value, ".")
	if !found {
		return nil, false
	}
	data, err := base64.RawURLEncoding.DecodeString(body)
	if err != nil {
		return nil, false
	}
	mac, err := base64.RawURLEncoding.DecodeString(sig)
	if err != nil {
		return nil, false
	}
	if !hmac.Equal(mac, s.sign(body)) {
		return nil, false
	}
	return data, true
}

func (s *Store) sign(body string) []byte {
	mac := hmac.New(sha256.New, s.secret)
	mac.Write([]byte(body))
	return mac.Sum(nil)
}

// Save re-signs and sets the cookie when the session changed.
func (s *Store) Save(w http.ResponseWriter, sess *Session) {
	if !sess.dirty {
		return
	}
	sess.dirty = false
	body, err := json.Marshal(payload{Data: sess.data, Iat: time.Now().Unix()})
	if err != nil {
		slog.Warn("session: cannot encode payload", "error", err)
		return
	}
	encoded := base64.RawURLEncoding.EncodeToString(body) +
		"." + base64.RawURLEncoding.EncodeToString(s.sign(base64.RawURLEncoding.EncodeToString(body)))
	http.SetCookie(w, &http.Cookie{
		Name:     CookieName,
		Value:    encoded,
		Path:     "/",
		MaxAge:   int(s.ttl.Seconds()),
		HttpOnly: true,
		Secure:   s.secure,
		SameSite: http.SameSiteStrictMode,
	})
}

// Destroy expires the session cookie.
func (s *Store) Destroy(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     CookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		Secure:   s.secure,
		SameSite: http.SameSiteStrictMode,
	})
}

// Middleware loads the session before the handler and persists it
// afterwards when it changed. The response header is not flushed until
// the first body write (or the end of the request), so the session
// cookie survives handlers that only set a status code.
func (s *Store) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sess := s.Load(r)
		recorder := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(recorder, r.WithContext(WithContext(r.Context(), sess)))
		if sess.dirty {
			s.Save(recorder, sess)
		}
		recorder.flush()
	})
}

// WithContext returns a context carrying sess the way the middleware
// does. A chain that must not buffer the response, a large download,
// attaches the session through this instead: the middleware records the
// whole body so it can still set a cookie after the handler ran, which
// is exactly the memory a streamed response must not pay. A session
// attached this way is read-only in effect: nothing flushes its dirty
// flag, so mutations do not persist.
func WithContext(ctx context.Context, sess *Session) context.Context {
	return context.WithValue(ctx, contextKey{}, sess)
}

// FromContext returns the session attached to the request, or an empty
// detached session when the middleware is not installed.
func FromContext(ctx context.Context) *Session {
	if sess, ok := ctx.Value(contextKey{}).(*Session); ok {
		return sess
	}
	return &Session{data: map[string]string{}}
}

// statusRecorder buffers the response so the session cookie can be
// added after the handler ran: net/http snapshots headers at the first
// body byte otherwise.
type statusRecorder struct {
	http.ResponseWriter
	status      int
	wroteHeader bool
	sent        bool
	buf         bytes.Buffer
}

func (r *statusRecorder) WriteHeader(code int) {
	if !r.wroteHeader {
		r.status = code
		r.wroteHeader = true
	}
}

func (r *statusRecorder) Write(b []byte) (int, error) {
	r.wroteHeader = true
	return r.buf.Write(b)
}

// flush sends the recorded status and buffered body once.
func (r *statusRecorder) flush() {
	if r.sent {
		return
	}
	r.sent = true
	r.ResponseWriter.WriteHeader(r.status)
	if r.buf.Len() > 0 {
		_, _ = r.ResponseWriter.Write(r.buf.Bytes())
	}
}