package session

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

const (
	sessionKeyPrefix = "session:"
	metadataField    = "_meta"
	schemaVersion    = "1"
	userFieldPrefix  = "v:"

	scriptInvalid int64 = -1
	scriptMissing int64 = 0
	scriptOK      int64 = 1
)

var (
	ErrSessionClosed      = errors.New("session is closed")
	ErrSessionNotFound    = errors.New("session: not found")
	ErrInvalidSessionData = errors.New("session: invalid Redis hash data")
)

var createSessionScript = redis.NewScript(`
redis.call("HSET", KEYS[1], "_meta", "1")
redis.call("PEXPIRE", KEYS[1], ARGV[1])
return 1
`)

var readSessionScript = redis.NewScript(`
local keytype = redis.call("TYPE", KEYS[1]).ok
if keytype == "none" then
  return {0}
end
if keytype ~= "hash" then
  return {-1}
end
if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
  return {-1}
end
local values = redis.call("HGETALL", KEYS[1])
table.insert(values, 1, 1)
return values
`)

var setSessionFieldScript = redis.NewScript(`
local keytype = redis.call("TYPE", KEYS[1]).ok
if keytype == "none" then
  return 0
end
if keytype ~= "hash" then
  return -1
end
if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
  return -1
end
redis.call("HSET", KEYS[1], ARGV[1], ARGV[2])
redis.call("PEXPIRE", KEYS[1], ARGV[3])
return 1
`)

var deleteSessionFieldScript = redis.NewScript(`
local keytype = redis.call("TYPE", KEYS[1]).ok
if keytype == "none" then
  return 0
end
if keytype ~= "hash" then
  return -1
end
if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
  return -1
end
redis.call("HDEL", KEYS[1], ARGV[1])
redis.call("PEXPIRE", KEYS[1], ARGV[2])
return 1
`)

var refreshSessionScript = redis.NewScript(`
local keytype = redis.call("TYPE", KEYS[1]).ok
if keytype == "none" then
  return 0
end
if keytype ~= "hash" then
  return -1
end
if redis.call("HGET", KEYS[1], "_meta") ~= "1" then
  return -1
end
redis.call("PEXPIRE", KEYS[1], ARGV[1])
return 1
`)

type SessionInterface interface {
	UpdateTTL() error
	Set(key string, value any) error
	Get(key string) any
	ID() string
	Delete(string) error
	Destroy() error
}

type Manager struct {
	redisDB           *redis.Client
	privateCookieName string
	maxLifeTTL        time.Duration
	secure            bool

	cacheMu sync.Mutex
	cache   map[string]*Session
}

type Session struct {
	sid     string
	storage map[string]json.RawMessage
	manager *Manager

	mutex               sync.RWMutex
	closed              bool
	initializationError error
	expiresAt           time.Time
	expiryTimer         *time.Timer
}

func NewManager(client *redis.Client, cookieName string, ttl time.Duration, secure bool) *Manager {
	if ttl < time.Second {
		panic("session: ttl must be at least 1 second")
	}
	if err := (&http.Cookie{Name: cookieName}).Valid(); err != nil {
		panic("session: invalid cookie name")
	}

	return &Manager{
		redisDB:           client,
		privateCookieName: cookieName,
		maxLifeTTL:        ttl,
		secure:            secure,
		cache:             make(map[string]*Session),
	}
}

func (manager *Manager) findOrCacheSession(sid string) *Session {
	manager.cacheMu.Lock()
	defer manager.cacheMu.Unlock()
	if existing, ok := manager.cache[sid]; ok {
		return existing
	}
	session := &Session{sid: sid, manager: manager}
	manager.cache[sid] = session
	return session
}

func (manager *Manager) removeSessionIfEqual(sid string, expected *Session) {
	manager.cacheMu.Lock()
	defer manager.cacheMu.Unlock()
	if manager.cache[sid] == expected {
		delete(manager.cache, sid)
	}
}

func (manager *Manager) createSession(ctx context.Context) (*Session, *http.Cookie, error) {
	sid, err := manager.sessionID()
	if err != nil {
		return nil, nil, err
	}

	newSession := &Session{sid: sid, manager: manager}
	manager.cacheMu.Lock()
	manager.cache[sid] = newSession
	manager.cacheMu.Unlock()

	newSession.mutex.Lock()
	_, err = manager.runStatusScriptContext(ctx, createSessionScript, sid, manager.maxLifeTTL.Milliseconds())
	if err == nil {
		newSession.storage = make(map[string]json.RawMessage)
		newSession.resetExpiryLocked(manager.maxLifeTTL)
		newSession.mutex.Unlock()
		return newSession, manager.newCookie(newSession.sid), nil
	}

	newSession.markClosedLocked()
	newSession.mutex.Unlock()
	manager.removeSessionIfEqual(sid, newSession)
	return nil, nil, err
}

func (manager *Manager) newCookie(sid string) *http.Cookie {
	return &http.Cookie{
		Name:     manager.privateCookieName,
		Value:    sid,
		Path:     "/",
		HttpOnly: true,
		Secure:   manager.secure,
		MaxAge:   int(manager.maxLifeTTL / time.Second),
		SameSite: http.SameSiteStrictMode,
	}
}

func (manager *Manager) CreateSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) {
	ctx, cancel := context.WithTimeout(r.Context(), time.Second)
	defer cancel()

	session, cookie, err := manager.createSession(ctx)
	if err != nil {
		return nil, err
	}
	http.SetCookie(w, cookie)
	return session, nil
}

func (manager *Manager) GetSession(w http.ResponseWriter, r *http.Request) (SessionInterface, error) {
	cookie, err := r.Cookie(manager.privateCookieName)
	if errors.Is(err, http.ErrNoCookie) || (err == nil && cookie.Value == "") {
		return nil, ErrSessionNotFound
	}
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithTimeout(r.Context(), time.Second)
	defer cancel()

	sid := cookie.Value
	session := manager.findOrCacheSession(sid)

	err = session.redisReadAndRefresh(ctx)
	if err == nil {
		http.SetCookie(w, manager.newCookie(session.sid))
		return session, nil
	}
	if session.isClosed() {
		manager.removeSessionIfEqual(sid, session)
	}
	if errors.Is(err, redis.Nil) || errors.Is(err, ErrSessionClosed) {
		return nil, ErrSessionNotFound
	}
	return nil, err
}

func (manager *Manager) sessionID() (string, error) {
	b := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

func (manager *Manager) runStatusScript(script *redis.Script, sid string, args ...any) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	return manager.runStatusScriptContext(ctx, script, sid, args...)
}

func (manager *Manager) runStatusScriptContext(ctx context.Context, script *redis.Script, sid string, args ...any) (int64, error) {
	return script.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid}, args...).Int64()
}

func (manager *Manager) runReadScript(ctx context.Context, sid string) (int64, map[string]json.RawMessage, error) {
	return parseReadResult(readSessionScript.Run(ctx, manager.redisDB, []string{sessionKeyPrefix + sid}))
}

func parseReadResult(cmd *redis.Cmd) (int64, map[string]json.RawMessage, error) {
	values, err := cmd.Slice()
	if err != nil {
		return 0, nil, err
	}

	status := values[0].(int64)
	if status != scriptOK {
		return status, nil, nil
	}

	storage := make(map[string]json.RawMessage, (len(values)-3)/2)
	for i := 1; i < len(values); i += 2 {
		field := values[i].(string)
		value := values[i+1].(string)
		if field == metadataField {
			continue
		}
		key, err := decodeUserField(field)
		if err != nil || !json.Valid([]byte(value)) {
			return 0, nil, ErrInvalidSessionData
		}
		storage[key] = json.RawMessage([]byte(value))
	}
	return status, storage, nil
}

func encodeUserField(key string) string {
	return userFieldPrefix + base64.RawURLEncoding.EncodeToString([]byte(key))
}

func decodeUserField(field string) (string, error) {
	if len(field) < len(userFieldPrefix) || field[:len(userFieldPrefix)] != userFieldPrefix {
		return "", ErrInvalidSessionData
	}
	encoded := field[len(userFieldPrefix):]
	decoded, err := base64.RawURLEncoding.DecodeString(encoded)
	if err != nil || base64.RawURLEncoding.EncodeToString(decoded) != encoded {
		return "", ErrInvalidSessionData
	}
	return string(decoded), nil
}

func (session *Session) redisReadAndRefresh(ctx context.Context) error {
	session.mutex.Lock()
	defer session.mutex.Unlock()
	if session.closed {
		if session.initializationError != nil {
			return session.initializationError
		}
		return ErrSessionClosed
	}

	wasLoaded := session.storage != nil
	status, values, err := session.manager.runReadScript(ctx, session.sid)
	if err != nil {
		if !wasLoaded {
			session.initializationError = err
			session.markClosedLocked()
		}
		return err
	}
	if status == scriptOK {
		status, err = session.manager.runStatusScriptContext(ctx, refreshSessionScript, session.sid, session.manager.maxLifeTTL.Milliseconds())
		if err != nil {
			if !wasLoaded {
				session.initializationError = err
				session.markClosedLocked()
			}
			return err
		}
	}
	if status == scriptMissing {
		session.markClosedLocked()
		return redis.Nil
	}
	if status == scriptOK {
		session.storage = values
		session.resetExpiryLocked(session.manager.maxLifeTTL)
		return nil
	}
	if !wasLoaded {
		session.initializationError = ErrInvalidSessionData
		session.markClosedLocked()
	}
	return ErrInvalidSessionData
}

func (session *Session) UpdateTTL() error {
	session.mutex.Lock()
	if session.closed {
		session.mutex.Unlock()
		return ErrSessionClosed
	}
	status, err := session.manager.runStatusScript(refreshSessionScript, session.sid, session.manager.maxLifeTTL.Milliseconds())
	if err != nil {
		session.mutex.Unlock()
		return err
	}
	switch status {
	case scriptOK:
		session.resetExpiryLocked(session.manager.maxLifeTTL)
		session.mutex.Unlock()
		return nil
	case scriptMissing:
		session.markClosedLocked()
		session.mutex.Unlock()
		session.manager.removeSessionIfEqual(session.sid, session)
		return redis.Nil
	default:
		session.mutex.Unlock()
		return ErrInvalidSessionData
	}
}

func (session *Session) ID() string {
	return session.sid
}

func (session *Session) Get(key string) any {
	session.mutex.RLock()
	if session.closed {
		session.mutex.RUnlock()
		return nil
	}
	raw, ok := session.storage[key]
	if !ok {
		session.mutex.RUnlock()
		return nil
	}
	value := append(json.RawMessage(nil), raw...)
	session.mutex.RUnlock()

	var decoded any
	if err := json.Unmarshal(value, &decoded); err != nil {
		return nil
	}
	return decoded
}

func (session *Session) Set(key string, value any) error {
	payload, err := json.Marshal(value)
	if err != nil {
		return err
	}

	session.mutex.Lock()
	if session.closed {
		session.mutex.Unlock()
		return ErrSessionClosed
	}
	status, err := session.manager.runStatusScript(setSessionFieldScript, session.sid, encodeUserField(key), string(payload), session.manager.maxLifeTTL.Milliseconds())
	if err != nil {
		session.mutex.Unlock()
		return err
	}
	switch status {
	case scriptOK:
		session.storage[key] = json.RawMessage(payload)
		session.resetExpiryLocked(session.manager.maxLifeTTL)
		session.mutex.Unlock()
		return nil
	case scriptMissing:
		session.markClosedLocked()
		session.mutex.Unlock()
		session.manager.removeSessionIfEqual(session.sid, session)
		return ErrSessionClosed
	default:
		session.mutex.Unlock()
		return ErrInvalidSessionData
	}
}

func (session *Session) Delete(key string) error {
	session.mutex.Lock()
	if session.closed {
		session.mutex.Unlock()
		return ErrSessionClosed
	}
	status, err := session.manager.runStatusScript(deleteSessionFieldScript, session.sid, encodeUserField(key), session.manager.maxLifeTTL.Milliseconds())
	if err != nil {
		session.mutex.Unlock()
		return err
	}
	switch status {
	case scriptOK:
		delete(session.storage, key)
		session.resetExpiryLocked(session.manager.maxLifeTTL)
		session.mutex.Unlock()
		return nil
	case scriptMissing:
		session.markClosedLocked()
		session.mutex.Unlock()
		session.manager.removeSessionIfEqual(session.sid, session)
		return ErrSessionClosed
	default:
		session.mutex.Unlock()
		return ErrInvalidSessionData
	}
}

func (session *Session) Destroy() error {
	session.mutex.Lock()
	if session.closed {
		session.mutex.Unlock()
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	err := session.manager.redisDB.Del(ctx, sessionKeyPrefix+session.sid).Err()
	cancel()
	if err != nil {
		session.mutex.Unlock()
		return err
	}
	session.markClosedLocked()
	session.mutex.Unlock()
	session.manager.removeSessionIfEqual(session.sid, session)
	return nil
}

func (session *Session) isClosed() bool {
	session.mutex.RLock()
	defer session.mutex.RUnlock()
	return session.closed
}

func (session *Session) markClosedLocked() {
	session.closed = true
	session.storage = nil
	if session.expiryTimer != nil {
		session.expiryTimer.Stop()
		session.expiryTimer = nil
	}
}

func (session *Session) resetExpiryLocked(ttl time.Duration) {
	session.expiresAt = time.Now().Add(ttl)
	if session.expiryTimer != nil {
		session.expiryTimer.Stop()
	}
	session.expiryTimer = time.AfterFunc(ttl, session.expire)
}

func (session *Session) expire() {
	session.mutex.Lock()
	if session.closed || time.Now().Before(session.expiresAt) {
		session.mutex.Unlock()
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	pttl, err := session.manager.redisDB.PTTL(ctx, sessionKeyPrefix+session.sid).Result()
	cancel()
	if err != nil {
		session.resetExpiryLocked(min(session.manager.maxLifeTTL, time.Minute))
		session.mutex.Unlock()
		return
	}
	if pttl >= 0 {
		session.resetExpiryLocked(pttl + time.Millisecond)
		session.mutex.Unlock()
		return
	}
	session.markClosedLocked()
	session.mutex.Unlock()
	session.manager.removeSessionIfEqual(session.sid, session)
}