package store

import (
	"context"

	"github.com/filecoin-project/lotus/chain/types"
	"github.com/filecoin-project/specs-actors/actors/abi"
	lru "github.com/hashicorp/golang-lru"
	"golang.org/x/xerrors"
)

type ChainIndex struct {
	skipCache *lru.ARCCache

	loadTipSet loadTipSetFunc

	skipLength abi.ChainEpoch
}
type loadTipSetFunc func(types.TipSetKey) (*types.TipSet, error)

func NewChainIndex(lts loadTipSetFunc) *ChainIndex {
	sc, _ := lru.NewARC(8192)
	return &ChainIndex{
		skipCache:  sc,
		loadTipSet: lts,
		skipLength: 20,
	}
}

type lbEntry struct {
	ts           *types.TipSet
	parentHeight abi.ChainEpoch
	target       types.TipSetKey
}

func (ci *ChainIndex) GetTipsetByHeight(ctx context.Context, from *types.TipSet, to abi.ChainEpoch) (*types.TipSet, error) {
	if from.Height()-to <= ci.skipLength {
		return ci.walkBack(from, to)
	}

	rounded, err := ci.roundDown(from)
	if err != nil {
		return nil, err
	}

	cur := rounded.Key()
	for {
		cval, ok := ci.skipCache.Get(cur)
		if !ok {
			fc, err := ci.fillCache(cur)
			if err != nil {
				return nil, err
			}
			cval = fc
		}

		lbe := cval.(*lbEntry)
		if lbe.ts.Height() == to || lbe.parentHeight < to {
			return lbe.ts, nil
		} else if to > lbe.ts.Height()-ci.skipLength {
			return ci.walkBack(from, to)
		}

		cur = lbe.target
	}
}

func (ci *ChainIndex) fillCache(tsk types.TipSetKey) (*lbEntry, error) {
	ts, err := ci.loadTipSet(tsk)
	if err != nil {
		return nil, err
	}

	// will either be equal to ts.Height, or at least > ts.Parent.Height()
	rheight := ci.roundHeight(ts.Height())

	parent, err := ci.loadTipSet(ts.Parents())
	if err != nil {
		return nil, err
	}

	if parent.Height() > rheight {
		return nil, xerrors.Errorf("cache is inconsistent")
	}

	rheight -= ci.skipLength

	skipTarget, err := ci.walkBack(parent, rheight)
	if err != nil {
		return nil, err
	}

	lbe := &lbEntry{
		ts:           ts,
		parentHeight: parent.Height(),
		target:       skipTarget.Key(),
	}
	ci.skipCache.Add(tsk, lbe)

	return lbe, nil
}

func (ci *ChainIndex) roundHeight(h abi.ChainEpoch) abi.ChainEpoch {
	return abi.ChainEpoch(h/ci.skipLength) * ci.skipLength
}

func (ci *ChainIndex) roundDown(ts *types.TipSet) (*types.TipSet, error) {
	target := ci.roundHeight(ts.Height())

	rounded, err := ci.walkBack(ts, target)
	if err != nil {
		return nil, err
	}

	return rounded, nil
}

func (ci *ChainIndex) walkBack(from *types.TipSet, to abi.ChainEpoch) (*types.TipSet, error) {
	if to > from.Height() {
		return nil, xerrors.Errorf("looking for tipset with height greater than start point")
	}

	if to == from.Height() {
		return from, nil
	}

	ts := from

	for {
		pts, err := ci.loadTipSet(ts.Parents())
		if err != nil {
			return nil, err
		}

		if to > pts.Height() {
			return ts, nil
		}
		if to == pts.Height() {
			return pts, nil
		}

		ts = pts
	}
}