package tree

import (
	"errors"
	"sync"
	"time"
)

type (
	circuitBreaker struct {
		breakDuration time.Duration
		threshold     int

		mu    sync.RWMutex
		state map[uint64]state
	}

	state struct {
		counter        int
		breakTimestamp time.Time
	}
)

var ErrCBClosed = errors.New("circuit breaker is closed")

func newCircuitBreaker(breakDuration time.Duration, threshold int) *circuitBreaker {
	return &circuitBreaker{
		breakDuration: breakDuration,
		threshold:     threshold,
		state:         make(map[uint64]state),
	}
}

func (cb *circuitBreaker) checkBreak(id uint64) error {
	cb.mu.RLock()
	s, ok := cb.state[id]
	cb.mu.RUnlock()

	if ok && time.Since(s.breakTimestamp) < cb.breakDuration {
		return ErrCBClosed
	}

	return nil
}

func (cb *circuitBreaker) openBreak(id uint64) {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	delete(cb.state, id)
}

func (cb *circuitBreaker) incError(id uint64) {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	s := cb.state[id]

	s.counter++
	if s.counter >= cb.threshold {
		s.counter = cb.threshold
		if time.Since(s.breakTimestamp) >= cb.breakDuration {
			s.breakTimestamp = time.Now()
		}
	}

	cb.state[id] = s
}

func (cb *circuitBreaker) Do(id uint64, f func() error) error {
	if err := cb.checkBreak(id); err != nil {
		return err
	}

	err := f()
	if err == nil {
		cb.openBreak(id)
	} else {
		cb.incError(id)
	}

	return err
}