package staking

import (
	"encoding/gob"
	"os"
	"sync"
)

type StakeInfo struct {
	Amount   uint64
	Duration uint64
}

type StakingStore struct {
	data map[string]StakeInfo
	mu   sync.RWMutex
}

func NewStakingStore() *StakingStore {
	return &StakingStore{
		data: make(map[string]StakeInfo),
	}
}

func (s *StakingStore) Stake(address string, amount uint64, duration uint64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data[address] = StakeInfo{Amount: amount, Duration: duration}
	return nil
}

func (s *StakingStore) Unstake(address string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.data, address)
	return nil
}

func (s *StakingStore) GetStakeInfo(address string) (StakeInfo, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	info, ok := s.data[address]
	return info, ok
}

func (s *StakingStore) UpdateStakeInfo(address string, info StakeInfo) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data[address] = info
}

func (s *StakingStore) Snapshot() map[string]StakeInfo {
	s.mu.RLock()
	defer s.mu.RUnlock()
	clone := make(map[string]StakeInfo)
	for k, v := range s.data {
		clone[k] = v
	}
	return clone
}

func (s *StakingStore) SaveToDisk(filename string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	encoder := gob.NewEncoder(file)
	return encoder.Encode(s.data)
}

func (s *StakingStore) LoadFromDisk(filename string) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	decoder := gob.NewDecoder(file)
	return decoder.Decode(&s.data)
}