kodiak.go

package protocols

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"

	bind "github.com/ethereum/go-ethereum/accounts/abi/bind/v2"
	"github.com/ethereum/go-ethereum/common"
	"github.com/infrared-dao/protocols/fetchers"
	"github.com/infrared-dao/protocols/internal/sc"
	"github.com/infrared-dao/protocols/multicall3"
	"github.com/rs/zerolog"
	"github.com/shopspring/decimal"
)

// Compile-time check that KodiakLPPriceProvider satisfies the
// BatchablePriceProvider extension interface, so a missing PriceReads or
// ComputePrice method fails to compile rather than silently degrading to
// the legacy fallback path at runtime.
var _ BatchablePriceProvider = &KodiakLPPriceProvider{}

const (
	V2PoolByteCodeSHA256    = "bffda5c9e111fa411890cc93d58cb5c58a14f97bcbfa642efe300c766c4397f6"
	V3IslandByteCodeSHA256  = "f1b68b8044f372c1831075be05a63c1757ef9c64cabd5dfbf9a749a93f25b1da"
	CharmPoolByteCodeSHA256 = "b8de812c5ad434d1d101b9dfb46de284637741a9b36b25e856fe1eca8a1388e0"
)

// enforce interface adherence
var _ Protocol = &KodiakLPPriceProvider{}
var _ KodiakContract = &KodiakV3Island{}
var _ KodiakContract = &KodiakV2Pool{}
var _ KodiakContract = &KodiakCharmPool{}

type Balances struct {
	Amount0 *big.Int
	Amount1 *big.Int
}

// Wrapper types which implement a common KodiakContract interface
type KodiakV3Island struct {
	*sc.KodiakIsland
}
type KodiakV2Pool struct {
	*sc.UniswapV2
}
type KodiakCharmPool struct {
	*sc.AlphaProVault
}

type KodiakContract interface {
	Decimals(opts *bind.CallOpts) (uint8, error)
	Token0(opts *bind.CallOpts) (common.Address, error)
	Token1(opts *bind.CallOpts) (common.Address, error)
	TotalSupply(opts *bind.CallOpts) (*big.Int, error)
	GetBalances(opts *bind.CallOpts) (Balances, error)

	// Multicall3 support — each variant packs/unpacks its own balances
	// method (V2 calls getReserves, V3 calls getUnderlyingBalances, Charm
	// calls getTotalAmounts0) using its underlying generated-binding ABI.

	// PackBalancesCalldata returns the ABI-encoded calldata for the
	// variant-specific balances method, suitable for use as a Multicall3
	// Call3.CallData.
	PackBalancesCalldata() ([]byte, error)
	// UnpackBalancesResponse decodes the raw return data from the
	// variant-specific balances method into the unified Balances type.
	UnpackBalancesResponse(data []byte) (Balances, error)
}

// Unify different implementations for getting internal balance amounts

func (v3 *KodiakV3Island) GetBalances(opts *bind.CallOpts) (Balances, error) {
	balances, err := v3.GetUnderlyingBalances(opts)
	if err != nil {
		return Balances{}, err
	}
	return Balances{
		Amount0: balances.Amount0Current,
		Amount1: balances.Amount1Current,
	}, nil
}

func (v3 *KodiakV3Island) PackBalancesCalldata() ([]byte, error) {
	a, err := sc.KodiakIslandMetaData.GetAbi()
	if err != nil {
		return nil, fmt.Errorf("KodiakV3Island: get ABI: %w", err)
	}
	return a.Pack("getUnderlyingBalances")
}

func (v3 *KodiakV3Island) UnpackBalancesResponse(data []byte) (Balances, error) {
	a, err := sc.KodiakIslandMetaData.GetAbi()
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakV3Island: get ABI: %w", err)
	}
	out, err := a.Methods["getUnderlyingBalances"].Outputs.Unpack(data)
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakV3Island: unpack getUnderlyingBalances: %w", err)
	}
	// getUnderlyingBalances returns (amount0Current, amount1Current) plus
	// additional accumulator fields we don't need; take the first two.
	if len(out) < 2 {
		return Balances{}, fmt.Errorf("KodiakV3Island: getUnderlyingBalances returned %d values, want >=2", len(out))
	}
	a0, ok := out[0].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakV3Island: amount0Current type %T, want *big.Int", out[0])
	}
	a1, ok := out[1].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakV3Island: amount1Current type %T, want *big.Int", out[1])
	}
	return Balances{Amount0: a0, Amount1: a1}, nil
}

func (v2 *KodiakV2Pool) GetBalances(opts *bind.CallOpts) (Balances, error) {
	balances, err := v2.GetReserves(opts)
	if err != nil {
		return Balances{}, err
	}
	return Balances{
		Amount0: balances.Reserve0,
		Amount1: balances.Reserve1,
	}, nil
}

func (v2 *KodiakV2Pool) PackBalancesCalldata() ([]byte, error) {
	a, err := sc.UniswapV2MetaData.GetAbi()
	if err != nil {
		return nil, fmt.Errorf("KodiakV2Pool: get ABI: %w", err)
	}
	return a.Pack("getReserves")
}

func (v2 *KodiakV2Pool) UnpackBalancesResponse(data []byte) (Balances, error) {
	a, err := sc.UniswapV2MetaData.GetAbi()
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakV2Pool: get ABI: %w", err)
	}
	out, err := a.Methods["getReserves"].Outputs.Unpack(data)
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakV2Pool: unpack getReserves: %w", err)
	}
	// getReserves returns (reserve0 uint112, reserve1 uint112,
	// blockTimestampLast uint32). uint112 unpacks as *big.Int.
	if len(out) < 2 {
		return Balances{}, fmt.Errorf("KodiakV2Pool: getReserves returned %d values, want >=2", len(out))
	}
	r0, ok := out[0].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakV2Pool: reserve0 type %T, want *big.Int", out[0])
	}
	r1, ok := out[1].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakV2Pool: reserve1 type %T, want *big.Int", out[1])
	}
	return Balances{Amount0: r0, Amount1: r1}, nil
}

func (kc *KodiakCharmPool) GetBalances(opts *bind.CallOpts) (Balances, error) {
	balances, err := kc.GetTotalAmounts0(opts)
	if err != nil {
		return Balances{}, err
	}
	return Balances{
		Amount0: balances.Total0,
		Amount1: balances.Total1,
	}, nil
}

func (kc *KodiakCharmPool) PackBalancesCalldata() ([]byte, error) {
	a, err := sc.AlphaProVaultMetaData.GetAbi()
	if err != nil {
		return nil, fmt.Errorf("KodiakCharmPool: get ABI: %w", err)
	}
	return a.Pack("getTotalAmounts0")
}

func (kc *KodiakCharmPool) UnpackBalancesResponse(data []byte) (Balances, error) {
	a, err := sc.AlphaProVaultMetaData.GetAbi()
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakCharmPool: get ABI: %w", err)
	}
	out, err := a.Methods["getTotalAmounts0"].Outputs.Unpack(data)
	if err != nil {
		return Balances{}, fmt.Errorf("KodiakCharmPool: unpack getTotalAmounts0: %w", err)
	}
	if len(out) < 2 {
		return Balances{}, fmt.Errorf("KodiakCharmPool: getTotalAmounts0 returned %d values, want >=2", len(out))
	}
	t0, ok := out[0].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakCharmPool: total0 type %T, want *big.Int", out[0])
	}
	t1, ok := out[1].(*big.Int)
	if !ok {
		return Balances{}, fmt.Errorf("KodiakCharmPool: total1 type %T, want *big.Int", out[1])
	}
	return Balances{Amount0: t0, Amount1: t1}, nil
}

// Define core types for the Kodiak Adapter

type KodiakConfig struct {
	Token0      string `json:"token0"`
	Token1      string `json:"token1"`
	LPTDecimals uint   `json:"lpt_decimals"`
}

// KodiakLPPriceProvider defines the provider for Kodiak LP price and TVL.
type KodiakLPPriceProvider struct {
	address     common.Address
	block       *big.Int
	priceMap    map[string]Price
	logger      zerolog.Logger
	configBytes []byte
	config      *KodiakConfig
	contract    KodiakContract
}

// NewKodiakLPPriceProvider creates a new instance of the KodiakLPPriceProvider.
func NewKodiakLPPriceProvider(
	address common.Address,
	block *big.Int,
	prices map[string]Price,
	logger zerolog.Logger,
	config []byte,
) *KodiakLPPriceProvider {
	k := &KodiakLPPriceProvider{
		address:     address,
		block:       block,
		priceMap:    prices,
		logger:      logger,
		configBytes: config,
	}
	return k
}

// Initialize checks the configuration/data provided and instantiates the KodiakV1 smart contract.
func (k *KodiakLPPriceProvider) Initialize(ctx context.Context, client bind.ContractBackend, _ fetchers.HttpClient) error {
	var err error

	k.config = &KodiakConfig{}
	err = json.Unmarshal(k.configBytes, k.config)
	if err != nil {
		k.logger.Error().Err(err).Msg("failed to deserialize config")
		return err
	}
	_, ok := k.priceMap[k.config.Token0]
	if !ok {
		err = fmt.Errorf("no price data found for token0 (%s)", k.config.Token0)
		k.logger.Error().Msg(err.Error())
		return err
	}
	_, ok = k.priceMap[k.config.Token1]
	if !ok {
		err = fmt.Errorf("no price data found for token1 (%s)", k.config.Token1)
		k.logger.Error().Msg(err.Error())
		return err
	}

	// initialize Kodiak contract of correct type
	k.contract, err = newKodiakContract(ctx, client, k.address)
	if err != nil {
		k.logger.Error().Err(err).Msg("failed to instantiate Kodiak smart contract")
		return err
	}

	return nil
}

// LPTokenPrice returns the current price of LP token in USD.
//
// Legacy path — issues two sequential eth_calls (totalSupply + balances).
// New batchable path (PriceReads + ComputePrice) is preferred when the
// caller supports it; both paths share computeLPPriceFromReads so prices
// match exactly between paths for the same inputs.
func (k *KodiakLPPriceProvider) LPTokenPrice(ctx context.Context) (string, error) {
	// Fetch total supply
	totalSupply, err := k.getTotalSupply(ctx)
	if err != nil {
		return "", err
	}

	// Fetch balances
	amount0, amount1, err := k.getBalances(ctx)
	if err != nil {
		return "", err
	}

	price, err := k.computeLPPriceFromReads(totalSupply, Balances{Amount0: amount0, Amount1: amount1})
	if err != nil {
		return "", err
	}
	return price.StringFixed(roundingDecimals), nil
}

// PriceReads describes the two eth_calls needed to compute the LP token
// price: totalSupply (standard ERC20) + the variant-specific balances
// method. Returned in fixed order so ComputePrice can index by position.
func (k *KodiakLPPriceProvider) PriceReads() ([]multicall3.Call3, error) {
	totalSupplyData, err := kodiakTotalSupplyCalldata()
	if err != nil {
		return nil, fmt.Errorf("kodiak: pack totalSupply: %w", err)
	}
	balancesData, err := k.contract.PackBalancesCalldata()
	if err != nil {
		return nil, fmt.Errorf("kodiak: pack balances: %w", err)
	}
	return []multicall3.Call3{
		{Target: k.address, AllowFailure: true, CallData: totalSupplyData},
		{Target: k.address, AllowFailure: true, CallData: balancesData},
	}, nil
}

// ComputePrice decodes the responses from PriceReads (in the same order)
// and computes the LP token price. Pure compute — no network I/O.
func (k *KodiakLPPriceProvider) ComputePrice(responses []multicall3.Result3) (decimal.Decimal, error) {
	if len(responses) != 2 {
		return decimal.Zero, fmt.Errorf("kodiak: expected 2 responses, got %d", len(responses))
	}
	if !responses[0].Success {
		return decimal.Zero, errors.New("kodiak: totalSupply call reverted in multicall")
	}
	if !responses[1].Success {
		return decimal.Zero, errors.New("kodiak: balances call reverted in multicall")
	}

	totalSupply, err := unpackKodiakTotalSupply(responses[0].ReturnData)
	if err != nil {
		return decimal.Zero, fmt.Errorf("kodiak: unpack totalSupply: %w", err)
	}
	balances, err := k.contract.UnpackBalancesResponse(responses[1].ReturnData)
	if err != nil {
		return decimal.Zero, fmt.Errorf("kodiak: unpack balances: %w", err)
	}

	return k.computeLPPriceFromReads(totalSupply, balances)
}

// computeLPPriceFromReads is the shared math between the legacy
// LPTokenPrice and the new batchable ComputePrice paths. Pure: takes the
// raw read values and config/prices, returns the LP token price.
func (k *KodiakLPPriceProvider) computeLPPriceFromReads(totalSupply *big.Int, balances Balances) (decimal.Decimal, error) {
	if totalSupply.Sign() == 0 {
		err := errors.New("totalSupply is zero, cannot calculate LP token price")
		k.logger.Error().Err(err).Msg("Invalid totalSupply")
		return decimal.Zero, err
	}

	price0, err := k.getPrice(k.config.Token0)
	if err != nil {
		return decimal.Zero, err
	}
	price1, err := k.getPrice(k.config.Token1)
	if err != nil {
		return decimal.Zero, err
	}

	amount0Decimal := NormalizeAmount(balances.Amount0, price0.Decimals)
	amount1Decimal := NormalizeAmount(balances.Amount1, price1.Decimals)
	totalValue := amount0Decimal.Mul(price0.Price).Add(amount1Decimal.Mul(price1.Price))

	totalSupplyDecimal := NormalizeAmount(totalSupply, k.config.LPTDecimals)
	pricePerToken := totalValue.Div(totalSupplyDecimal)

	k.logger.Debug().
		Str("totalValue", totalValue.String()).
		Str("totalSupply", totalSupplyDecimal.String()).
		Str("pricePerToken", pricePerToken.String()).
		Msg("LP token price calculated successfully")

	return pricePerToken, nil
}

// kodiakTotalSupplyCalldata returns the ABI-encoded calldata for ERC20
// totalSupply(). Standard signature shared by every KodiakContract variant.
func kodiakTotalSupplyCalldata() ([]byte, error) {
	a, err := sc.KodiakIslandMetaData.GetAbi()
	if err != nil {
		return nil, err
	}
	return a.Pack("totalSupply")
}

// unpackKodiakTotalSupply decodes a totalSupply response into *big.Int.
// Standard ERC20 return shape; reused across all KodiakContract variants.
func unpackKodiakTotalSupply(data []byte) (*big.Int, error) {
	a, err := sc.KodiakIslandMetaData.GetAbi()
	if err != nil {
		return nil, err
	}
	out, err := a.Methods["totalSupply"].Outputs.Unpack(data)
	if err != nil {
		return nil, err
	}
	if len(out) != 1 {
		return nil, fmt.Errorf("totalSupply returned %d values, want 1", len(out))
	}
	v, ok := out[0].(*big.Int)
	if !ok {
		return nil, fmt.Errorf("totalSupply value type %T, want *big.Int", out[0])
	}
	return v, nil
}

// TVL returns the Total Value Locked in the pool/island as USD
func (k *KodiakLPPriceProvider) TVL(ctx context.Context) (string, error) {
	totalValue, err := k.totalValue(ctx)
	if err != nil {
		return "", err
	}

	k.logger.Debug().
		Str("totalValue", totalValue.String()).
		Msg("TVL calculated successfully")

	return totalValue.StringFixed(roundingDecimals), nil
}

func (k *KodiakLPPriceProvider) GetConfig(ctx context.Context, address string, client bind.ContractBackend) ([]byte, error) {
	var err error
	if !common.IsHexAddress(address) {
		err = fmt.Errorf("invalid smart contract address, '%s'", address)
		return nil, err
	}

	// initialize Kodiak contract of correct type
	contract, err := newKodiakContract(ctx, client, common.HexToAddress(address))
	if err != nil {
		k.logger.Error().Err(err).Msg("failed to instantiate Kodiak smart contract")
		return nil, err
	}

	kc := KodiakConfig{}
	opts := &bind.CallOpts{
		Context: ctx,
	}

	// token0
	addr, err := contract.Token0(opts)
	if err != nil {
		err = fmt.Errorf("failed to obtain token0 address for kodiak vault %s, %v", address, err)
		return nil, err
	}
	kc.Token0 = strings.ToLower(addr.Hex())

	// token1
	addr, err = contract.Token1(opts)
	if err != nil {
		err = fmt.Errorf("failed to obtain token1 address for kodiak vault %s, %v", address, err)
		return nil, err
	}
	kc.Token1 = strings.ToLower(addr.Hex())

	// decimals
	decimals, err := contract.Decimals(opts)
	if err != nil {
		err = fmt.Errorf("failed to obtain number of decimals for LP token %s, %v", address, err)
		return nil, err
	}
	kc.LPTDecimals = uint(decimals)

	body, err := json.Marshal(kc)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func (k *KodiakLPPriceProvider) UpdateBlock(block *big.Int, prices map[string]Price) {
	k.block = block
	if prices != nil {
		k.priceMap = prices
	}
}

// TVLBreakdown returns the breakdown of TVL by underlying tokens (token0 and token1).
func (k *KodiakLPPriceProvider) TVLBreakdown(ctx context.Context) (map[string]TokenTVL, error) {
	// Fetch internal balances
	amount0, amount1, err := k.getBalances(ctx)
	if err != nil {
		return nil, err
	}

	price0, err := k.getPrice(k.config.Token0)
	if err != nil {
		return nil, err
	}
	amount0Decimal := NormalizeAmount(amount0, price0.Decimals)
	token0USDValue := amount0Decimal.Mul(price0.Price)

	price1, err := k.getPrice(k.config.Token1)
	if err != nil {
		return nil, err
	}
	amount1Decimal := NormalizeAmount(amount1, price1.Decimals)
	token1USDValue := amount1Decimal.Mul(price1.Price)

	totalValue := token0USDValue.Add(token1USDValue)

	// Calculate ratios (handle zero TVL case)
	var token0Ratio, token1Ratio decimal.Decimal
	if totalValue.IsZero() {
		token0Ratio = decimal.Zero
		token1Ratio = decimal.Zero
	} else {
		token0Ratio = token0USDValue.Div(totalValue)
		token1Ratio = token1USDValue.Div(totalValue)
	}

	breakdown := map[string]TokenTVL{
		k.config.Token0: {
			TokenAddress: k.config.Token0,
			TokenSymbol:  price0.TokenName,
			Amount:       amount0Decimal,
			USDValue:     token0USDValue,
			Ratio:        token0Ratio,
		},
		k.config.Token1: {
			TokenAddress: k.config.Token1,
			TokenSymbol:  price1.TokenName,
			Amount:       amount1Decimal,
			USDValue:     token1USDValue,
			Ratio:        token1Ratio,
		},
	}

	k.logger.Debug().
		Str("token0", k.config.Token0).
		Str("token0Amount", amount0Decimal.String()).
		Str("token0USDValue", token0USDValue.String()).
		Str("token0Ratio", token0Ratio.String()).
		Str("token1", k.config.Token1).
		Str("token1Amount", amount1Decimal.String()).
		Str("token1USDValue", token1USDValue.String()).
		Str("token1Ratio", token1Ratio.String()).
		Msg("TVL breakdown calculated successfully")

	return breakdown, nil
}

// Internal Helper methods not able to be called except in this file

// totalValue is the legacy USD-value helper used by TVL() and TVLBreakdown().
// LPTokenPrice now uses computeLPPriceFromReads directly so legacy and
// batchable paths share identical math.
func (k *KodiakLPPriceProvider) totalValue(ctx context.Context) (decimal.Decimal, error) {
	var err error

	// Fetch internal balances
	amount0, amount1, err := k.getBalances(ctx)
	if err != nil {
		return decimal.Zero, err
	}

	price0, err := k.getPrice(k.config.Token0)
	if err != nil {
		return decimal.Zero, err
	}
	amount0Decimal := NormalizeAmount(amount0, price0.Decimals)

	price1, err := k.getPrice(k.config.Token1)
	if err != nil {
		return decimal.Zero, err
	}
	amount1Decimal := NormalizeAmount(amount1, price1.Decimals)
	totalValue := amount0Decimal.Mul(price0.Price).Add(amount1Decimal.Mul(price1.Price))
	return totalValue, nil
}

func (k *KodiakLPPriceProvider) getPrice(tokenKey string) (*Price, error) {
	price, ok := k.priceMap[tokenKey]
	if !ok {
		err := fmt.Errorf("no price data found for token (%s)", tokenKey)
		k.logger.Error().Msg(err.Error())
		return nil, err
	}
	return &price, nil
}

// getTotalSupply fetches the total supply of the LP token.
func (k *KodiakLPPriceProvider) getTotalSupply(ctx context.Context) (*big.Int, error) {
	opts := &bind.CallOpts{
		Context:     ctx,
		BlockNumber: k.block,
	}

	ts, err := k.contract.TotalSupply(opts)
	if err != nil {
		k.logger.Error().Msgf("failed to obtain total supply for kodiak contract %s, %v",
			k.address.String(), err)
		return nil, fmt.Errorf("failed to get kodiak contract total supply, err: %w", err)
	}

	return ts, err
}

// getUnderlyingBalances fetches the underlying token balances.
func (k *KodiakLPPriceProvider) getBalances(ctx context.Context) (*big.Int, *big.Int, error) {
	opts := &bind.CallOpts{
		Context:     ctx,
		BlockNumber: k.block,
	}

	bs, err := k.contract.GetBalances(opts)
	if err != nil {
		k.logger.Error().Msgf("failed to obtain balances for kodiak contract %s, %v",
			k.address.String(), err)
		return nil, nil, fmt.Errorf("failed to get kodiak contract balances, err: %w", err)
	}
	return bs.Amount0, bs.Amount1, err
}

// Check type of contract and create new connection
func newKodiakContract(
	ctx context.Context,
	client bind.ContractBackend,
	address common.Address,
) (KodiakContract, error) {
	contractHash := checkContractCodeHash(ctx, client, address)

	switch contractHash {
	case V2PoolByteCodeSHA256:
		// initialize contract as a Kodiak V2 Pool
		pool, err := sc.NewUniswapV2(address, client)
		if err != nil {
			return nil, err
		}
		return &KodiakV2Pool{pool}, nil
	case V3IslandByteCodeSHA256:
		// initialize contract as a Kodiak V3 Island
		island, err := sc.NewKodiakIsland(address, client)
		if err != nil {
			return nil, err
		}
		return &KodiakV3Island{island}, nil
	case CharmPoolByteCodeSHA256:
		// initialize contract as a Kodiak x Charm Pool (concentrated liquidity)
		pool, err := sc.NewAlphaProVault(address, client)
		if err != nil {
			return nil, err
		}
		return &KodiakCharmPool{pool}, nil
	default:
		// Error for yet unknown contract type being attempted
		return nil, errors.New("contract unrecognized: neither V2 Pool or V3 Island")
	}
}

// get contract byte code and SHA256 hash it for checks on the contract type against known values
func checkContractCodeHash(
	ctx context.Context,
	client bind.ContractBackend,
	address common.Address,
) string {
	bytecode, err := client.CodeAt(ctx, address, nil)
	if err != nil {
		panic(err)
	}
	byteCodeString := hex.EncodeToString(bytecode)

	hash := sha256.New()
	hash.Write([]byte(byteCodeString))
	return hex.EncodeToString(hash.Sum(nil))
}