package v1

import (
	"time"

	"gno.land/p/nt/avl"
	"gno.land/p/nt/ufmt"

	i256 "gno.land/p/gnoswap/int256"
	u256 "gno.land/p/gnoswap/uint256"

	"gno.land/r/gnoswap/common"
	sr "gno.land/r/gnoswap/staker"
)

var q128 = u256.MustFromDecimal("340282366920938463463374607431768211456")

// Pools represents the global pool storage
type Pools struct {
	tree *avl.Tree // string poolPath -> pool
}

func NewPools() *Pools {
	return &Pools{
		tree: avl.NewTree(),
	}
}

// Get returns the pool for the given poolPath
func (self *Pools) Get(poolPath string) (*sr.Pool, bool) {
	v, ok := self.tree.Get(poolPath)
	if !ok {
		return nil, false
	}
	p, ok := v.(*sr.Pool)
	if !ok {
		panic(ufmt.Sprintf("failed to cast v to *Pool: %T", v))
	}
	return p, true
}

// GetPoolOrNil returns the pool for the given poolPath, or returns nil if it does not exist
func (self *Pools) GetPoolOrNil(poolPath string) *sr.Pool {
	pool, ok := self.Get(poolPath)
	if !ok {
		return nil
	}
	return pool
}

// set sets the pool for the given poolPath
func (self *Pools) set(poolPath string, pool *sr.Pool) {
	self.tree.Set(poolPath, pool)
}

// Has returns true if the pool exists for the given poolPath
func (self *Pools) Has(poolPath string) bool {
	return self.tree.Has(poolPath)
}

func (self *Pools) IterateAll(fn func(key string, pool *sr.Pool) bool) {
	self.tree.Iterate("", "", func(key string, value any) bool {
		p, ok := value.(*sr.Pool)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to *Pool: %T", value))
		}
		return fn(key, p)
	})
}

type PoolResolver struct {
	*sr.Pool
}

func (self *PoolResolver) IncentivesResolver() *IncentivesResolver {
	return NewIncentivesResolver(self.Incentives())
}

// Get the latest global reward ratio accumulation in [0, currentTime] range.
// Returns the time and the accumulation.
func (self *PoolResolver) CurrentGlobalRewardRatioAccumulation(currentTime int64) (time int64, acc *u256.Uint) {
	acc = u256.Zero()

	self.GlobalRewardRatioAccumulation().ReverseIterate(0, currentTime, func(key int64, value any) bool {
		time = key
		v, ok := value.(*u256.Uint)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to *u256.Uint: %T", value))
		}
		acc = v
		return true
	})
	if acc == nil {
		panic("should not happen, globalRewardRatioAccumulation must be set when pool is created")
	}
	return time, acc
}

// Get the latest tick in [0, currentTime] range.
// Returns the tick.
func (self *PoolResolver) CurrentTick(currentTime int64) (tick int32) {
	self.HistoricalTick().ReverseIterate(0, currentTime, func(key int64, value any) bool {
		res, ok := value.(int32)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to int32: %T", value))
		}
		tick = res
		return true
	})
	return tick
}

func (self *PoolResolver) CurrentStakedLiquidity(currentTime int64) (liquidity *u256.Uint) {
	liquidity = u256.Zero()

	self.StakedLiquidity().ReverseIterate(0, currentTime, func(key int64, value any) bool {
		res, ok := value.(*u256.Uint)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to *u256.Uint: %T", value))
		}
		liquidity = res
		return true
	})
	return liquidity
}

func (self *PoolResolver) TickResolver(tickId int32) *TickResolver {
	return NewTickResolver(self.Ticks().Get(tickId))
}

// IsExternallyIncentivizedPool returns true if the pool has any active external incentives.
func (self *PoolResolver) IsExternallyIncentivizedPool() bool {
	currentTime := time.Now().Unix()
	hasIncentive := false
	self.Incentives().IncentiveTrees().Iterate("", "", func(key string, value any) bool {
		incentive, ok := value.(*sr.ExternalIncentive)
		if !ok {
			panic("failed to cast value to *ExternalIncentive")
		}

		resolver := NewExternalIncentiveResolver(incentive)
		if !resolver.IsEnded(currentTime) {
			hasIncentive = true
			return true
		}

		return false
	})

	return hasIncentive
}

// Get the latest reward in [0, currentTime] range.
// Returns the reward.
func (self *PoolResolver) CurrentReward(currentTime int64) (reward int64) {
	self.RewardCache().ReverseIterate(0, currentTime, func(key int64, value any) bool {
		res, ok := value.(int64)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to int64: %T", value))
		}
		reward = res
		return true
	})
	return reward
}

// cacheReward sets the current reward for the pool
// If the pool is in unclaimable period, it will end the unclaimable period, updates the reward, and start the unclaimable period again.
//
// Important behavior for initial tier assignment:
// - When a pool first receives a tier, oldTierReward=0 and currentTierReward>0
// - If the pool has zero liquidity at this point, startUnclaimablePeriod() is called
// - This ensures unclaimable period tracking begins from the moment rewards start emitting
func (self *PoolResolver) cacheReward(currentTime int64, currentTierReward int64) {
	oldTierReward := self.CurrentReward(currentTime)
	if oldTierReward == currentTierReward {
		return
	}

	isInUnclaimable := self.CurrentStakedLiquidity(currentTime).IsZero()
	if isInUnclaimable {
		// End any existing unclaimable period
		// Note: If lastUnclaimableTime is 0 (not yet tracking), this is a no-op
		self.endUnclaimablePeriod(currentTime)
	}

	self.RewardCache().Set(currentTime, currentTierReward)

	if isInUnclaimable {
		// Start/restart unclaimable period tracking
		// This handles initial tier assignment when lastUnclaimableTime is 0
		self.startUnclaimablePeriod(currentTime)
	}
}

// cacheInternalReward caches the current emission and updates the global reward ratio accumulation.
func (self *PoolResolver) cacheInternalReward(currentTime int64, currentEmission int64) {
	self.cacheReward(currentTime, currentEmission)

	currentStakedLiquidity := self.CurrentStakedLiquidity(currentTime)
	self.updateGlobalRewardRatioAccumulation(currentTime, currentStakedLiquidity)
}

func (self *PoolResolver) calculateGlobalRewardRatioAccumulation(currentTime int64, currentStakedLiquidity *u256.Uint) *u256.Uint {
	oldAccTime, oldAcc := self.CurrentGlobalRewardRatioAccumulation(currentTime)
	timeDiff := safeSubInt64(currentTime, oldAccTime)
	if timeDiff == 0 {
		return oldAcc.Clone()
	}
	if timeDiff < 0 {
		panic("time cannot go backwards")
	}

	if currentStakedLiquidity.IsZero() {
		return oldAcc.Clone()
	}

	acc := u256.MulDiv(
		u256.NewUintFromInt64(timeDiff),
		q128,
		currentStakedLiquidity,
	)
	return u256.Zero().Add(oldAcc, acc)
}

// updateGlobalRewardRatioAccumulation updates the global reward ratio accumulation and returns the new accumulation.
func (self *PoolResolver) updateGlobalRewardRatioAccumulation(currentTime int64, currentStakedLiquidity *u256.Uint) *u256.Uint {
	newAcc := self.calculateGlobalRewardRatioAccumulation(currentTime, currentStakedLiquidity)

	self.GlobalRewardRatioAccumulation().Set(currentTime, newAcc)
	return newAcc
}

// RewardStateOf initializes a new RewardState for the given deposit.
func (self *PoolResolver) RewardStateOf(deposit *sr.Deposit) *RewardState {
	warmups := len(deposit.Warmups())
	result := &RewardState{
		pool:      self,
		deposit:   NewDepositResolver(deposit),
		rewards:   make([]int64, warmups),
		penalties: make([]int64, warmups),
	}

	for i := range result.rewards {
		result.rewards[i] = 0
		result.penalties[i] = 0
	}

	return result
}

// reset clears cached rewards/penalties so a RewardState can be reused without re-allocating.
func (self *RewardState) reset() {
	for i := range self.rewards {
		self.rewards[i] = 0
		self.penalties[i] = 0
	}
}

// NewPool creates a new pool with the given poolPath and currentHeight.
func NewPoolResolver(pool *sr.Pool) *PoolResolver {
	return &PoolResolver{
		Pool: pool,
	}
}

// RewardState is a struct for storing the intermediate state for reward calculation.
type RewardState struct {
	pool    *PoolResolver
	deposit *DepositResolver

	// accumulated rewards for each warmup
	rewards   []int64
	penalties []int64
}

// calculateInternalReward calculates the internal reward for the deposit.
// It calls rewardPerWarmup for each rewardCache interval, applies warmup, and returns the rewards and penalties.
func (self *RewardState) calculateInternalReward(startTime, endTime int64) ([]int64, []int64) {
	currentReward := self.pool.CurrentReward(startTime)

	self.pool.RewardCache().Iterate(startTime, endTime, func(key int64, value any) bool {
		reward, ok := value.(int64)
		if !ok {
			panic(ufmt.Sprintf("failed to cast value to int64: %T", value))
		}

		// Calculate reward for the period before this cache entry
		err := self.rewardPerWarmup(startTime, key, currentReward)
		if err != nil {
			panic(err)
		}

		startTime = key
		currentReward = reward
		return false
	})

	if startTime < endTime {
		// For the remaining period, use the last cached reward value
		// If the pool was de-tiered and the last cached value is 0, this ensures no new rewards
		err := self.rewardPerWarmup(startTime, endTime, currentReward)
		if err != nil {
			panic(err)
		}
	}

	self.applyWarmup()

	return self.rewards, self.penalties
}

// updateExternalReward updates the external reward for the deposit.
// It updates the last collect time for the external reward for the given incentive ID.
// It returns an error if the current time is less than the last collect time for the external reward for the given incentive ID.
func (self *RewardState) updateExternalReward(startTime, endTime int64, incentive *sr.ExternalIncentive) error {
	lastCollectTime := self.deposit.ExternalRewardLastCollectTime(incentive.IncentiveId())
	if startTime < lastCollectTime {
		// This must not happen, but adding some guards just in case.
		startTime = lastCollectTime
	}

	ictvStart := incentive.StartTimestamp()
	if endTime < ictvStart {
		return nil // Not started yet
	}

	if startTime < ictvStart {
		startTime = ictvStart
	}

	ictvEnd := incentive.EndTimestamp()
	if endTime > ictvEnd {
		endTime = ictvEnd
	}

	if startTime > ictvEnd {
		return nil // Already ended
	}

	return self.rewardPerWarmup(startTime, endTime, incentive.RewardPerSecond())
}

// calculateCollectableExternalReward calculates the calculated external reward for the deposit.
// It calls updateExternalReward for the incentive period, applies warmup and returns the rewards and penalties.
// used for reward calculation for a calculatable incentive
func (self *RewardState) calculateCollectableExternalReward(startTime, endTime int64, incentive *sr.ExternalIncentive) int64 {
	err := self.updateExternalReward(startTime, endTime, incentive)
	if err != nil {
		panic(err)
	}

	currentReward := u256.Zero()

	for i := range self.rewards {
		currentReward = currentReward.Add(currentReward, u256.NewUintFromInt64(self.rewards[i]))
	}

	return safeConvertToInt64(currentReward)
}

// calculateExternalReward calculates the external reward for the deposit.
// It calls rewardPerWarmup for startTime to endTime(clamped to the incentive period), applies warmup and returns the rewards and penalties.
func (self *RewardState) calculateExternalReward(startTime, endTime int64, incentive *sr.ExternalIncentive) ([]int64, []int64) {
	err := self.updateExternalReward(startTime, endTime, incentive)
	if err != nil {
		panic(err)
	}

	// apply warmup to collect rewards
	self.applyWarmup()

	return self.rewards, self.penalties
}

// applyWarmup applies the warmup to the rewards and calculate penalties.
func (self *RewardState) applyWarmup() {
	for i, warmup := range self.deposit.Warmups() {
		warmupReward := self.rewards[i]

		// calculate warmup reward applying warmup ratio
		self.rewards[i] = safeMulInt64(warmupReward, int64(warmup.WarmupRatio)) / 100

		// warmup penalty is the difference between the warmup reward and the warmup reward applying warmup ratio
		self.penalties[i] = safeSubInt64(warmupReward, self.rewards[i])
	}
}

// rewardPerWarmup calculates the reward for each warmup, adds to the RewardState's rewards array.
func (self *RewardState) rewardPerWarmup(startTime, endTime int64, rewardPerSecond int64) error {
	// Return early if startTime equals endTime to avoid unnecessary computation
	if startTime == endTime {
		return nil
	}

	for i, warmup := range self.deposit.Warmups() {
		if startTime >= warmup.NextWarmupTime {
			// passed the warmup
			continue
		}

		if endTime < warmup.NextWarmupTime {
			rewardAcc := self.pool.CalculateRewardForPosition(
				startTime,
				self.pool.CurrentTick(startTime),
				endTime,
				self.pool.CurrentTick(endTime),
				self.deposit.Deposit,
			)

			rewardAcc = u256.Zero().Mul(rewardAcc, self.deposit.Liquidity())
			rewardAcc = u256.MulDiv(rewardAcc, u256.NewUintFromInt64(rewardPerSecond), q128)
			self.rewards[i] = safeAddInt64(self.rewards[i], safeConvertToInt64(rewardAcc))

			break
		}

		rewardAcc := self.pool.CalculateRewardForPosition(
			startTime,
			self.pool.CurrentTick(startTime),
			warmup.NextWarmupTime,
			self.pool.CurrentTick(warmup.NextWarmupTime),
			self.deposit.Deposit,
		)

		rewardAcc = u256.Zero().Mul(rewardAcc, self.deposit.Liquidity())
		rewardAcc = u256.MulDiv(rewardAcc, u256.NewUintFromInt64(rewardPerSecond), q128)
		self.rewards[i] = safeAddInt64(self.rewards[i], safeConvertToInt64(rewardAcc))

		startTime = warmup.NextWarmupTime
	}

	return nil
}

// modifyDeposit updates the pool's staked liquidity and returns the new staked liquidity.
// updates when there is a change in the staked liquidity(tick cross, stake, unstake)
func (self *PoolResolver) modifyDeposit(delta *i256.Int, currentTime int64, nextTick int32) *u256.Uint {
	// update staker side pool info
	lastStakedLiquidity := self.CurrentStakedLiquidity(currentTime)
	deltaApplied := common.LiquidityMathAddDelta(lastStakedLiquidity, delta)
	result := self.updateGlobalRewardRatioAccumulation(currentTime, lastStakedLiquidity)

	// historical tick does NOT actually reflect the tick at the timestamp, but it provides correct ordering for the staked positions
	// because TickCrossHook is assured to be called for the staked-initialized ticks
	self.HistoricalTick().Set(currentTime, nextTick)

	switch deltaApplied.Sign() {
	case -1:
		panic("stakedLiquidity is less than 0, should not happen")
	case 0:
		if lastStakedLiquidity.Sign() == 1 {
			// StakedLiquidity moved from positive to zero, start unclaimable period
			self.startUnclaimablePeriod(currentTime)
			self.IncentivesResolver().startUnclaimablePeriod(currentTime)
		}
	case 1:
		if lastStakedLiquidity.Sign() == 0 {
			// StakedLiquidity moved from zero to positive, end unclaimable period
			self.endUnclaimablePeriod(currentTime)
			self.IncentivesResolver().endUnclaimablePeriod(currentTime)
		}
	}

	self.Pool.StakedLiquidity().Set(currentTime, deltaApplied)

	return result
}

// startUnclaimablePeriod starts the unclaimable period.
func (self *PoolResolver) startUnclaimablePeriod(currentTime int64) {
	if self.LastUnclaimableTime() == 0 {
		// We set only if it's the first time entering(0 indicates not set yet)
		self.SetLastUnclaimableTime(currentTime)
	}
}

// endUnclaimablePeriod ends the unclaimable period.
// Accumulates to unclaimableAcc and resets lastUnclaimableTime to 0.
func (self *PoolResolver) endUnclaimablePeriod(currentTime int64) {
	if self.LastUnclaimableTime() == 0 {
		// lastUnclaimableTime = 0 means tracking hasn't started yet
		// This is normal during initial pool creation or when called from cacheReward
		// during tier assignment with zero liquidity
		return
	}

	self.updateUnclaimableAccumulateRewards(currentTime)
	self.SetLastUnclaimableTime(0)
}

// updateUnclaimableAccumulateRewards ends the unclaimable period.
// Accumulates to unclaimableAcc and resets lastUnclaimableTime to 0.
func (self *PoolResolver) updateUnclaimableAccumulateRewards(currentTime int64) {
	if self.LastUnclaimableTime() >= currentTime {
		return
	}

	unclaimableDuration := safeSubInt64(currentTime, self.LastUnclaimableTime())
	currentUnclaimableReward := safeMulInt64(unclaimableDuration, self.CurrentReward(self.LastUnclaimableTime()))
	self.SetUnclaimableAcc(safeAddInt64(self.UnclaimableAcc(), currentUnclaimableReward))
}

// processUnclaimableReward processes the unclaimable reward and returns the accumulated reward.
// It resets unclaimableAcc to 0 and properly manages lastUnclaimableTime based on pool state.
func (self *PoolResolver) processUnclaimableReward(endTime int64) int64 {
	// Check current pool liquidity state
	isZeroStakedLiquidity := self.CurrentStakedLiquidity(endTime).IsZero()

	if self.LastUnclaimableTime() > 0 {
		// We have an ongoing unclaimable period tracking
		self.updateUnclaimableAccumulateRewards(endTime)

		if isZeroStakedLiquidity {
			// Still unclaimable - accumulate rewards up to endTime
			// Update tracking time for continuing unclaimable period
			self.SetLastUnclaimableTime(endTime)
		} else {
			// Was unclaimable but now has liquidity - properly end the period
			self.SetLastUnclaimableTime(0)
		}
	} else {
		if isZeroStakedLiquidity {
			// No previous tracking but currently unclaimable - this shouldn't normally happen
			// as startUnclaimablePeriod should have been called when liquidity reached 0
			// Start tracking from now
			self.SetLastUnclaimableTime(endTime)
		}
	}

	// Return and reset accumulated unclaimable rewards
	internalUnClaimable := self.UnclaimableAcc()
	self.SetUnclaimableAcc(0)
	return internalUnClaimable
}

// Calculates reward for a position *without* considering debt or warmup
// It calculates the theoretical total reward for the position if it has been staked since the pool creation
func (self *PoolResolver) CalculateRawRewardForPosition(currentTime int64, currentTick int32, deposit *sr.Deposit) *u256.Uint {
	var rewardAcc *u256.Uint

	globalAcc := self.calculateGlobalRewardRatioAccumulation(currentTime, self.CurrentStakedLiquidity(currentTime))

	lowerAcc := self.TickResolver(deposit.TickLower()).CurrentOutsideAccumulation(currentTime)
	upperAcc := self.TickResolver(deposit.TickUpper()).CurrentOutsideAccumulation(currentTime)
	if currentTick < deposit.TickLower() {
		rewardAcc = u256.Zero().Sub(lowerAcc, upperAcc)
	} else if currentTick >= deposit.TickUpper() {
		rewardAcc = u256.Zero().Sub(upperAcc, lowerAcc)
	} else {
		rewardAcc = u256.Zero().Sub(globalAcc, lowerAcc)
		rewardAcc = rewardAcc.Sub(rewardAcc, upperAcc)
	}

	return rewardAcc
}

// Calculate actual reward in [startTime, endTime) for a position by
// subtracting the startTime's raw reward from the endTime's raw reward
func (self *PoolResolver) CalculateRewardForPosition(
	startTime int64,
	startTick int32,
	endTime int64,
	endTick int32,
	deposit *sr.Deposit,
) *u256.Uint {
	rewardAcc := self.CalculateRawRewardForPosition(endTime, endTick, deposit)
	debtAcc := self.CalculateRawRewardForPosition(startTime, startTick, deposit)

	return u256.Zero().Sub(rewardAcc, debtAcc)
}