// Package auth implements shared-password authentication using a salted
// PBKDF2-HMAC-SHA256 password hash, stateless HMAC-signed session cookies,
// and an in-memory login rate limiter. It depends only on the standard library.
package auth

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	pbkdf2Iterations = 120_000
	hashKeyLen       = 32
	hashScheme       = "pbkdf2-sha256"
)

// Service handles authentication concerns for the gateway.
type Service struct {
	passwordHash []byte // verified against the scheme string in parseHash
	secret       []byte
	cookieName   string
	ttl          time.Duration
	limiter      *rateLimiter
}

// New creates an auth Service.
func New(passwordHash string, secret []byte, cookieName string, ttl time.Duration) *Service {
	return &Service{
		passwordHash: []byte(passwordHash),
		secret:       secret,
		cookieName:   cookieName,
		ttl:          ttl,
		limiter:      newRateLimiter(5, time.Minute),
	}
}

// CheckPassword verifies a plaintext password against the stored hash.
func (s *Service) CheckPassword(plain string) bool {
	if len(s.passwordHash) == 0 {
		return false
	}
	scheme, iter, salt, want, err := parseHash(string(s.passwordHash))
	if err != nil {
		return false
	}
	if scheme != hashScheme {
		return false
	}
	got := pbkdf2Key([]byte(plain), salt, iter, hashKeyLen)
	return subtle.ConstantTimeCompare(got, want) == 1
}

// HashPassword returns a self-describing salted hash for a plaintext password.
func HashPassword(plain string) (string, error) {
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return "", err
	}
	key := pbkdf2Key([]byte(plain), salt, pbkdf2Iterations, hashKeyLen)
	return hashScheme + "$" +
		strconv.Itoa(pbkdf2Iterations) + "$" +
		base64.RawStdEncoding.EncodeToString(salt) + "$" +
		base64.RawStdEncoding.EncodeToString(key), nil
}

func parseHash(h string) (scheme string, iter int, salt, key []byte, err error) {
	parts := strings.Split(h, "$")
	if len(parts) != 4 {
		return "", 0, nil, nil, errors.New("invalid hash format")
	}
	iter, err = strconv.Atoi(parts[1])
	if err != nil || iter <= 0 {
		return "", 0, nil, nil, errors.New("invalid iteration count")
	}
	salt, err = base64.RawStdEncoding.DecodeString(parts[2])
	if err != nil {
		return "", 0, nil, nil, err
	}
	key, err = base64.RawStdEncoding.DecodeString(parts[3])
	if err != nil {
		return "", 0, nil, nil, err
	}
	return parts[0], iter, salt, key, nil
}

// pbkdf2Key implements PBKDF2-HMAC-SHA256 (RFC 2898).
func pbkdf2Key(password, salt []byte, iter, keyLen int) []byte {
	prf := hmac.New(sha256.New, password)
	hLen := prf.Size()
	numBlocks := (keyLen + hLen - 1) / hLen

	out := make([]byte, 0, numBlocks*hLen)
	var block [4]byte
	for i := 1; i <= numBlocks; i++ {
		prf.Reset()
		prf.Write(salt)
		binary.BigEndian.PutUint32(block[:], uint32(i))
		prf.Write(block[:])
		u := prf.Sum(nil)
		t := make([]byte, len(u))
		copy(t, u)
		for j := 1; j < iter; j++ {
			prf.Reset()
			prf.Write(u)
			u = prf.Sum(u[:0])
			for k := range t {
				t[k] ^= u[k]
			}
		}
		out = append(out, t...)
	}
	return out[:keyLen]
}

// IssueSession creates a signed session token (cookie value).
// Format: <expUnix>.<base64url-hmac>.
func (s *Service) IssueSession(now time.Time) string {
	exp := now.Add(s.ttl).Unix()
	payload := strconv.FormatInt(exp, 10)
	mac := s.computeMAC(payload)
	return payload + "." + base64.RawURLEncoding.EncodeToString(mac)
}

// VerifySession validates a session token and returns true if valid & not expired.
func (s *Service) VerifySession(token string, now time.Time) bool {
	parts := strings.SplitN(token, ".", 2)
	if len(parts) != 2 {
		return false
	}
	payload := parts[0]
	macStr := parts[1]

	mac, err := base64.RawURLEncoding.DecodeString(macStr)
	if err != nil {
		return false
	}
	expected := s.computeMAC(payload)
	if !hmac.Equal(mac, expected) {
		return false
	}
	exp, err := strconv.ParseInt(payload, 10, 64)
	if err != nil {
		return false
	}
	return exp > now.Unix()
}

func (s *Service) computeMAC(payload string) []byte {
	m := hmac.New(sha256.New, s.secret)
	m.Write([]byte(payload))
	return m.Sum(nil)
}

// SetSessionCookie writes the session cookie on the response.
func (s *Service) SetSessionCookie(w http.ResponseWriter, token string) {
	http.SetCookie(w, &http.Cookie{
		Name:     s.cookieName,
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   int(s.ttl.Seconds()),
	})
}

// ClearSessionCookie expires the session cookie.
func (s *Service) ClearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     s.cookieName,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	})
}

// SessionFromRequest extracts and validates the session cookie.
func (s *Service) SessionFromRequest(r *http.Request, now time.Time) bool {
	c, err := r.Cookie(s.cookieName)
	if err != nil {
		return false
	}
	return s.VerifySession(c.Value, now)
}

// AllowLogin enforces a per-IP rate limit on login attempts.
func (s *Service) AllowLogin(ip string) bool {
	return s.limiter.allow(ip)
}

// CookieName returns the configured cookie name.
func (s *Service) CookieName() string { return s.cookieName }

// TTL returns the configured session TTL.
func (s *Service) TTL() time.Duration { return s.ttl }

// ---- rate limiter (fixed window per IP, in-memory) ----

type rateLimiter struct {
	mu     sync.Mutex
	max    int
	window time.Duration
	hits   map[string][]time.Time
}

func newRateLimiter(max int, window time.Duration) *rateLimiter {
	return &rateLimiter{max: max, window: window, hits: make(map[string][]time.Time)}
}

func (r *rateLimiter) allow(ip string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	now := time.Now()
	cutoff := now.Add(-r.window)
	fresh := r.hits[ip][:0]
	for _, t := range r.hits[ip] {
		if t.After(cutoff) {
			fresh = append(fresh, t)
		}
	}
	if len(fresh) >= r.max {
		r.hits[ip] = fresh
		return false
	}
	fresh = append(fresh, now)
	r.hits[ip] = fresh
	return true
}