package pump_parser

import (
	"encoding/binary"
	"fmt"

	"github.com/gagliardetto/solana-go"
	"github.com/shopspring/decimal"
)

func decodeRaydiumClmmSwapArgs(data []byte) (amountSpecified uint64, otherAmountThreshold uint64, swapMode SwapMode, err error) {
	if len(data) < 41 {
		return 0, 0, SwapModeUnknown, fmt.Errorf("raydium clmm swap instruction data too short")
	}
	amountSpecified = binary.LittleEndian.Uint64(data[8:16])
	otherAmountThreshold = binary.LittleEndian.Uint64(data[16:24])
	isBaseInput := data[40] != 0
	swapMode = SwapModeExactOut
	if isBaseInput {
		swapMode = SwapModeExactIn
	}
	return amountSpecified, otherAmountThreshold, swapMode, nil
}

func raydiumClmmParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	if !tx.rawTx.accountList[instruction.ProgramIDIndex].Equals(raydiumClmmProgramID) {
		return nil, increaseOffset(offset), fmt.Errorf("raydiumClmm instruction not found,  offset, %d, %d", offset[0], offset[1])
	}

	decode := instruction.Data
	if len(decode) < 8 {
		return nil, increaseOffset(offset), fmt.Errorf("raydiumClmm program instruction data too short, offset, %d, %d", offset[0], offset[1])
	}

	discriminator := *(*[8]byte)(decode[:8])

	switch discriminator {
	case raydiumClmmCreatePoolDiscriminator:
		return raydiumClmmCreatePoolParser(tx, instruction, innerInstructions, offset)
	case raydiumClmmIncreaseLiquidityDiscriminator, raydiumClmmIncreaseLiquidityV2Discriminator, raydiumClmmOpenPositionDiscriminator, raydiumClmmOpenPositionV2Discriminator, raydiumClmmOpenPositionWithToken22NftDiscriminator:
		return raydiumClmmAddLiquidityParser(tx, instruction, innerInstructions, offset)
	case raydiumClmmDecreaseLiquidityDiscriminator, raydiumClmmDecreaseLiquidityV2Discriminator:
		return raydiumClmmDecreaseLiquidityParser(tx, instruction, innerInstructions, offset)
	case raydiumClmmCollectFundFeeDiscriminator, raydiumClmmCollectProtocolFeeDiscriminator:
		return raydiumClmmCollectFeeParser(tx, instruction, innerInstructions, offset)
	case raydiumClmmSwapDiscriminator, raydiumClmmSwapV2Discriminator:
		return raydiumClmmSwapParser(tx, instruction, innerInstructions, offset)
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}
}

func raydiumClmmCreatePoolParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	if len(instruction.Accounts) < 13 {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for raydiumClmm create pool instruction, offset, %d, %d", offset[0], offset[1])
	}
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]

	pool := tx.rawTx.accountList[instruction.Accounts[2]]
	creator := tx.rawTx.accountList[instruction.Accounts[0]]

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, instruction.Accounts[5])
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token0 vault balance after tx: %v", err)
	}

	quoteTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, instruction.Accounts[6])
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token1 vault balance after tx: %v", err)
	}

	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)
	offset[1] += 9
	return []Swap{
		{
			Program:           SolProgramRaydiumCLMM,
			Event:             "create",
			Pool:              pool,
			BaseMint:          baseTokenBalance.MintAccount,
			QuoteMint:         quoteTokenBalance.MintAccount,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			Creator:           creator,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumClmmAddLiquidityParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	discriminator := *(*[8]byte)(instruction.Data[:8])
	var (
		accountMin int
		market     solana.PublicKey
		//token0     solana.PublicKey
		//token1     solana.PublicKey
		lpToken solana.PublicKey
		vault0  int
		vault1  int
	)
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]

	switch discriminator {
	case raydiumClmmIncreaseLiquidityDiscriminator:
		accountMin = 12
		market = tx.rawTx.accountList[instruction.Accounts[2]]
		vault0 = instruction.Accounts[9]
		vault1 = instruction.Accounts[10]
	case raydiumClmmIncreaseLiquidityV2Discriminator:
		accountMin = 15
		market = tx.rawTx.accountList[instruction.Accounts[2]]
		vault0 = instruction.Accounts[9]
		vault1 = instruction.Accounts[10]
		//token0 = tx.rawTx.accountList[instruction.Accounts[13]]
		//token1 = tx.rawTx.accountList[instruction.Accounts[14]]
	case raydiumClmmOpenPositionDiscriminator:
		accountMin = 19
		market = tx.rawTx.accountList[instruction.Accounts[5]]
		vault0 = instruction.Accounts[12]
		vault1 = instruction.Accounts[13]
		lpToken = tx.rawTx.accountList[instruction.Accounts[2]]
	case raydiumClmmOpenPositionV2Discriminator:
		accountMin = 22
		market = tx.rawTx.accountList[instruction.Accounts[5]]
		vault0 = instruction.Accounts[12]
		vault1 = instruction.Accounts[13]
		lpToken = tx.rawTx.accountList[instruction.Accounts[2]]
		//token0 = tx.rawTx.accountList[instruction.Accounts[20]]
		//token1 = tx.rawTx.accountList[instruction.Accounts[21]]
	case raydiumClmmOpenPositionWithToken22NftDiscriminator:
		accountMin = 20
		market = tx.rawTx.accountList[instruction.Accounts[4]]
		vault0 = instruction.Accounts[11]
		vault1 = instruction.Accounts[12]
		lpToken = tx.rawTx.accountList[instruction.Accounts[2]]
		//token0 = tx.rawTx.accountList[instruction.Accounts[18]]
		//token1 = tx.rawTx.accountList[instruction.Accounts[19]]
	default:
		return nil, increaseOffset(offset), fmt.Errorf("invalid discriminator")
	}

	if len(instruction.Accounts) < accountMin {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for raydiumClmm add liquidity instruction, offset, %d, %d", offset[0], offset[1])
	}

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault0)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token0 vault balance after tx: %v", err)
	}

	quoteTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault1)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token1 vault balance after tx: %v", err)
	}

	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)

	offset[1] += 2

	return []Swap{
		{
			Program:           SolProgramRaydiumCLMM,
			Event:             "add_liquidity",
			Pool:              market,
			BaseMint:          baseTokenBalance.MintAccount,
			QuoteMint:         quoteTokenBalance.MintAccount,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			LpMint:            lpToken,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumClmmDecreaseLiquidityParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	discriminator := *(*[8]byte)(instruction.Data[:8])
	var (
		accountMin int
		market     solana.PublicKey
		vault0     int
		vault1     int
	)

	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	if discriminator == raydiumClmmDecreaseLiquidityDiscriminator {
		accountMin = 14
	} else if discriminator == raydiumClmmDecreaseLiquidityV2Discriminator {
		accountMin = 16
	}
	if len(instruction.Accounts) < accountMin {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for decrease liquidity instruction")
	}
	market = tx.rawTx.accountList[instruction.Accounts[3]]
	vault0 = instruction.Accounts[5]
	vault1 = instruction.Accounts[6]
	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault0)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token0 vault balance after tx: %v", err)
	}

	quoteTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault1)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token1 vault balance after tx: %v", err)
	}

	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)

	offset[1] += 2

	return []Swap{
		{
			Program:           SolProgramRaydiumCLMM,
			Event:             "remove_liquidity",
			Pool:              market,
			BaseMint:          baseTokenBalance.MintAccount,
			QuoteMint:         quoteTokenBalance.MintAccount,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumClmmCollectFeeParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	if len(instruction.Accounts) < 11 {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for CollectFeeParser instruction")
	}
	pool := tx.rawTx.accountList[instruction.Accounts[1]]

	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	vault0 := instruction.Accounts[3]
	vault1 := instruction.Accounts[4]
	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault0)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token0 vault balance after tx: %v", err)
	}

	quoteTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, vault1)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get token1 vault balance after tx: %v", err)
	}

	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)

	offset[1] += 2

	return []Swap{
		{
			Program:           SolProgramRaydiumCLMM,
			Event:             "remove_liquidity",
			Pool:              pool,
			BaseMint:          baseTokenBalance.MintAccount,
			QuoteMint:         quoteTokenBalance.MintAccount,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumClmmSwapParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	discriminator := *(*[8]byte)(instruction.Data[:8])
	var (
		pool solana.PublicKey

		accountMin          int
		tokenInVault        int
		tokenOutVault       int
		userTokenInAccount  int
		userTokenOutAccount int
	)
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	amountSpecified, otherAmountThreshold, swapMode, err := decodeRaydiumClmmSwapArgs(instruction.Data)
	if err != nil {
		return nil, increaseOffset(offset), err
	}
	if discriminator == raydiumClmmSwapDiscriminator {
		accountMin = 9
		pool = tx.rawTx.accountList[instruction.Accounts[2]]
		userTokenInAccount = instruction.Accounts[3]
		userTokenOutAccount = instruction.Accounts[4]
		tokenInVault = instruction.Accounts[5]
		tokenOutVault = instruction.Accounts[6]
	} else if discriminator == raydiumClmmSwapV2Discriminator {
		accountMin = 13
		pool = tx.rawTx.accountList[instruction.Accounts[2]]
		userTokenInAccount = instruction.Accounts[3]
		userTokenOutAccount = instruction.Accounts[4]
		tokenInVault = instruction.Accounts[5]
		tokenOutVault = instruction.Accounts[6]
	}
	if len(instruction.Accounts) < accountMin {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for swap instruction")
	}

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, tokenInVault)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get tokenIn vault balance after tx: %w", err)
	}
	quoteTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, tokenOutVault)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get tokenOut vault balance after tx: %w", err)
	}
	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)
	baseTokenProgram := baseTokenBalance.ProgramIDAccount
	quoteTokenProgram := quoteTokenBalance.ProgramIDAccount
	baseMint := baseTokenBalance.MintAccount
	quoteMint := quoteTokenBalance.MintAccount
	baseMintDecimals := uint8(baseTokenBalance.UITokenAmount.Decimals)
	quoteMintDecimals := uint8(quoteTokenBalance.UITokenAmount.Decimals)

	userBase := getAccountBalanceAfterTx(tx.rawTx, userTokenInAccount)
	userQuote := getAccountBalanceAfterTx(tx.rawTx, userTokenOutAccount)

	inners, err := getInnerInstructions(innerInstructions, offset[1])
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get inner instructions: %w", err)
	}
	if len(inners) < 2 {
		return nil, increaseOffset(offset), fmt.Errorf("not enough inner instructions for swap instruction")
	}
	baseVaultAccount := tx.rawTx.accountList[tokenInVault]
	quoteVaultAccount := tx.rawTx.accountList[tokenOutVault]
	userBaseAccount := tx.rawTx.accountList[userTokenInAccount]
	userQuoteAccount := tx.rawTx.accountList[userTokenOutAccount]
	var baseAmount, quoteAmount decimal.Decimal
	var baseFound, quoteFound bool
	for i := 0; i < 2; i++ {
		inner := inners[i]
		from, to, amount, err := parseTokenTransfer(tx.rawTx, inner)
		if err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("failed to parse token transfer: %w", err)
		}
		if from.Equals(userBaseAccount) && to.Equals(baseVaultAccount) && !baseFound {
			baseAmount = decimal.NewFromUint64(amount)
			baseFound = true
		} else if from.Equals(quoteVaultAccount) && to.Equals(userQuoteAccount) && !quoteFound {
			quoteAmount = decimal.NewFromUint64(amount)
			quoteFound = true
		}
	}
	if !baseFound || !quoteFound {
		return nil, increaseOffset(offset), fmt.Errorf("failed to find token transfer in inner instructions")
	}

	offset[1] += 2

	swap := Swap{
		Program:           SolProgramRaydiumCLMM,
		Event:             "sell",
		Pool:              pool,
		BaseMint:          baseMint,
		QuoteMint:         quoteMint,
		BaseTokenProgram:  baseTokenProgram,
		QuoteTokenProgram: quoteTokenProgram,
		BaseMintDecimals:  baseMintDecimals,
		QuoteMintDecimals: quoteMintDecimals,
		User:              tx.rawTx.accountList[instruction.Accounts[0]],
		BaseAmount:        baseAmount,
		QuoteAmount:       quoteAmount,
		BaseReserve:       baseReserve,
		QuoteReserve:      quoteReserve,
		UserBaseBalance:   userBase,
		UserQuoteBalance:  userQuote,
		EntryContract:     entryContract,
	}
	swap.SetSwapAmountInfo(swapMode, decimal.NewFromUint64(amountSpecified), decimal.NewFromUint64(otherAmountThreshold))
	return []Swap{swap}, offset, nil

}