package ttlmap

import (
	"sync"
	"time"

	"gitee.com/MikeDDD/gu/guutil"
	"gitee.com/MikeDDD/gu/logs"
)

type entity struct {
	key      string
	value    any
	expireAt int64
}

type TtlMap struct {
	lock     *sync.RWMutex
	entities map[string]*entity
}

func New() *TtlMap {
	m := &TtlMap{
		lock:     new(sync.RWMutex),
		entities: make(map[string]*entity),
	}

	guutil.GoWithRecover(func() {
		for range time.NewTicker(time.Second).C {
			m.recover()
		}
	})

	return m
}

func (t *TtlMap) recover() {
	defer guutil.InDeferRecover()

	t.lock.Lock()
	defer t.lock.Unlock()

	for k, v := range t.entities {
		if time.Now().UnixMilli() > v.expireAt {
			logs.Debug("recover %v", k)
			delete(t.entities, k)
		}
	}
}

func (t *TtlMap) Set(key string, value any, ttl time.Duration) {
	t.lock.Lock()
	defer t.lock.Unlock()

	t.entities[key] = &entity{
		key:      key,
		value:    value,
		expireAt: time.Now().UnixMilli() + ttl.Milliseconds(),
	}
}

func (t *TtlMap) SetSpecific(key string, value any, expireAt int64) {
	t.lock.Lock()
	defer t.lock.Unlock()

	t.entities[key] = &entity{
		key:      key,
		value:    value,
		expireAt: expireAt,
	}
}

func (t *TtlMap) Get(key string) any {
	t.lock.RLock()
	defer t.lock.RUnlock()

	v, exists := t.entities[key]
	if !exists {
		return nil
	}

	if time.Now().UnixMilli() >= v.expireAt {
		return nil
	}

	return v.value
}