package model

import (
	"context"
	"sync"
	"time"
)

type MemoryStore struct {
	mu        sync.Mutex
	codes     map[string]codeEntry
	cooldowns map[string]time.Time
}

type codeEntry struct {
	code string
	exp  time.Time
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{codes: make(map[string]codeEntry), cooldowns: make(map[string]time.Time)}
}

func (m *MemoryStore) Put(_ context.Context, key, code string, ttlSeconds int) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.codes[key] = codeEntry{code: code, exp: time.Now().Add(time.Duration(ttlSeconds) * time.Second)}
	return nil
}

func (m *MemoryStore) Get(_ context.Context, key string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	e, ok := m.codes[key]
	if !ok || time.Now().After(e.exp) {
		delete(m.codes, key)
		return "", nil
	}
	return e.code, nil
}

func (m *MemoryStore) Del(_ context.Context, key string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.codes, key)
	delete(m.cooldowns, key)
	return nil
}

func (m *MemoryStore) SetCooldown(_ context.Context, key string, ttlSeconds int) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if exp, ok := m.cooldowns[key]; ok && time.Now().Before(exp) {
		return false, nil
	}
	m.cooldowns[key] = time.Now().Add(time.Duration(ttlSeconds) * time.Second)
	return true, nil
}