package main

import (
	"github.com/gogf/gf/frame/g"
	"github.com/gogf/gf/net/ghttp"
	"sync"
)

// MemoryStorage 是一个实现了 ghttp.Storage 接口的内存存储引擎
type MemoryStorage struct {
	data map[string]map[string]interface{}
	lock sync.RWMutex
}

// Set 用于设置 Session 数据
func (m *MemoryStorage) Set(sessionId string, key string, value interface{}) error {
	m.lock.Lock()
	defer m.lock.Unlock()
	if m.data == nil {
		m.data = make(map[string]map[string]interface{})
	}
	if _, ok := m.data[sessionId]; !ok {
		m.data[sessionId] = make(map[string]interface{})
	}
	m.data[sessionId][key] = value
	return nil
}

// Get 用于获取 Session 数据
func (m *MemoryStorage) Get(sessionId string, key string) interface{} {
	m.lock.RLock()
	defer m.lock.RUnlock()
	if sessionData, ok := m.data[sessionId]; ok {
		if value, ok := sessionData[key]; ok {
			return value
		}
	}
	return nil
}

// Remove 用于移除 Session 数据
func (m *MemoryStorage) Remove(sessionId string, key string) error {
	m.lock.Lock()
	defer m.lock.Unlock()
	if sessionData, ok := m.data[sessionId]; ok {
		delete(sessionData, key)
	}
	return nil
}

// NewMemoryStorage 创建一个新的 MemoryStorage 实例
func NewMemoryStorage() ghttp.Storage {
	return &MemoryStorage{
		data: make(map[string]map[string]interface{}),
	}
}

func main() {
	s := g.Server()

	// 设置 Session 存储引擎为自定义的 MemoryStorage
	s.SetConfigWithMap(g.Map{
		"session": g.Map{
			"storage": NewMemoryStorage(),
		},
	})

	s.BindHandler("/set", func(r *ghttp.Request) {
		r.Session.Set("key", "value")
		r.Response.Write("Session set success")
	})

	s.BindHandler("/get", func(r *ghttp.Request) {
		val := r.Session.Get("key")
		r.Response.Write("Session value:", val)
	})

	s.Run()
}