package pump_parser

import (
	"bytes"
	"fmt"

	agbinary "github.com/gagliardetto/binary"
	"github.com/shopspring/decimal"
)

type raydiumCPmmSwapBaseInputArgs struct {
	AmountIn         uint64
	MinimumAmountOut uint64
}

type raydiumCPmmSwapBaseOutputArgs struct {
	MaxAmountIn uint64
	AmountOut   uint64
}

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

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

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

	switch discriminator {
	case raydiumCPmmInitializeDiscriminator, raydiumCPmmInitializeWithPermissionDiscriminator:
		return raydiumCPmmCreatePoolParser(tx, instruction, innerInstructions, offset)
	case raydiumCPmmDepositDiscriminator:
		return raydiumCPmmDepositParser(tx, instruction, innerInstructions, offset)
	case raydiumCPmmWithdrawDiscriminator:
		return raydiumCPmmWithdrawParser(tx, instruction, innerInstructions, offset)
	case raydiumCPmmCollectProtocolFeeDiscriminator, raydiumCPmmCollectFundFeeDiscriminator:
		return raydiumCPmmCollectParser(tx, instruction, innerInstructions, offset)
	case raydiumCPmmSwapBaseInputDiscriminator, raydiumCPmmSwapBaseOutputDiscriminator:
		return raydiumCPmmSwapParser(tx, instruction, innerInstructions, offset)
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}
}

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

	pool := tx.rawTx.accountList[instruction.Accounts[3]]
	lpMint := tx.rawTx.accountList[instruction.Accounts[6]]
	creator := tx.rawTx.accountList[instruction.Accounts[0]]
	vault0 := instruction.Accounts[10]
	vault1 := instruction.Accounts[11]
	if bytes.Equal(instruction.Data[:8], raydiumCPmmInitializeWithPermissionDiscriminator[:]) {
		pool = tx.rawTx.accountList[instruction.Accounts[4]]
		lpMint = tx.rawTx.accountList[instruction.Accounts[7]]
		creator = tx.rawTx.accountList[instruction.Accounts[1]]
		vault0 = instruction.Accounts[11]
		vault1 = instruction.Accounts[12]
	}

	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] += 13
	return []Swap{
		{
			Program:           SolProgramRaydiumCPMM,
			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,
			LpMint:            lpMint,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, increaseOffset(offset), nil
}

func raydiumCPmmDepositParser(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 deposit instruction")
	}
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	market := tx.rawTx.accountList[instruction.Accounts[2]]
	token0User := tx.rawTx.accountList[instruction.Accounts[4]]
	token1User := tx.rawTx.accountList[instruction.Accounts[5]]

	token0Vault := tx.rawTx.accountList[instruction.Accounts[6]]
	token1Vault := tx.rawTx.accountList[instruction.Accounts[7]]

	token0 := tx.rawTx.accountList[instruction.Accounts[10]]
	token1 := tx.rawTx.accountList[instruction.Accounts[11]]
	lpTokenMint := tx.rawTx.accountList[instruction.Accounts[12]]

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, instruction.Accounts[6])
	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[7])
	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)

	inners, err := getInnerInstructions(innerInstructions, offset[1])

	var baseFound, quoteFound bool
	var baseAmount, quoteAmount decimal.Decimal
	for _, inner := range inners {
		from, to, amount, err := parseTokenTransfer(tx.rawTx, inner)
		if err != nil {
			continue
		}
		if from.Equals(token0User) && to.Equals(token0Vault) && !baseFound {
			baseFound = true
			baseAmount = decimal.NewFromUint64(amount)
		} else if from.Equals(token1User) && to.Equals(token1Vault) && !quoteFound {
			quoteFound = true
			quoteAmount = decimal.NewFromUint64(amount)
		}
		if baseFound && quoteFound {
			break
		}
	}
	if !baseFound || !quoteFound {
		return nil, increaseOffset(offset), fmt.Errorf("failed to find token transfer in inner instructions")
	}
	offset[1] += 3
	return []Swap{
		{
			Program:           SolProgramRaydiumCPMM,
			Event:             "add_liquidity",
			Pool:              market,
			BaseMint:          token0,
			QuoteMint:         token1,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			BaseAmount:        baseAmount,
			QuoteAmount:       quoteAmount,
			LpMint:            lpTokenMint,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumCPmmWithdrawParser(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 deposit instruction")
	}
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	market := tx.rawTx.accountList[instruction.Accounts[2]]
	token0User := tx.rawTx.accountList[instruction.Accounts[4]]
	token1User := tx.rawTx.accountList[instruction.Accounts[5]]

	token0Vault := tx.rawTx.accountList[instruction.Accounts[6]]
	token1Vault := tx.rawTx.accountList[instruction.Accounts[7]]

	token0 := tx.rawTx.accountList[instruction.Accounts[10]]
	token1 := tx.rawTx.accountList[instruction.Accounts[11]]
	lpTokenMint := tx.rawTx.accountList[instruction.Accounts[12]]

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, instruction.Accounts[6])
	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[7])
	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)

	inners, err := getInnerInstructions(innerInstructions, offset[1])

	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("failed to get inner instructions: %v", err)
	}
	var baseFound, quoteFound bool
	var baseAmount, quoteAmount decimal.Decimal
	for _, inner := range inners {
		from, to, amount, err := parseTokenTransfer(tx.rawTx, inner)
		if err != nil {
			continue
		}
		if to.Equals(token0User) && from.Equals(token0Vault) && !baseFound {
			baseFound = true
			baseAmount = decimal.NewFromUint64(amount)
		} else if to.Equals(token1User) && from.Equals(token1Vault) && !quoteFound {
			quoteFound = true
			quoteAmount = decimal.NewFromUint64(amount)
		}
		if baseFound && quoteFound {
			break
		}
	}
	if !baseFound || !quoteFound {
		return nil, increaseOffset(offset), fmt.Errorf("failed to find token transfer in inner instructions")
	}
	offset[1] += 3
	return []Swap{
		{
			Program:           SolProgramRaydiumCPMM,
			Event:             "remove_liquidity",
			Pool:              market,
			BaseMint:          token0,
			QuoteMint:         token1,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			BaseAmount:        baseAmount,
			QuoteAmount:       quoteAmount,
			LpMint:            lpTokenMint,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumCPmmCollectParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	if len(instruction.Accounts) < 12 {
		return nil, increaseOffset(offset), fmt.Errorf("not enough accounts for deposit instruction")
	}
	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	market := tx.rawTx.accountList[instruction.Accounts[2]]
	token0User := tx.rawTx.accountList[instruction.Accounts[8]]
	token1User := tx.rawTx.accountList[instruction.Accounts[9]]

	token0Vault := tx.rawTx.accountList[instruction.Accounts[4]]
	token1Vault := tx.rawTx.accountList[instruction.Accounts[5]]

	token0 := tx.rawTx.accountList[instruction.Accounts[6]]
	token1 := tx.rawTx.accountList[instruction.Accounts[7]]

	inners, err := getInnerInstructions(innerInstructions, offset[1])

	var baseFound, quoteFound bool
	var baseAmount, quoteAmount decimal.Decimal
	for _, inner := range inners {
		from, to, amount, err := parseTokenTransfer(tx.rawTx, inner)
		if err != nil {
			continue
		}
		if to.Equals(token0User) && from.Equals(token0Vault) && !baseFound {
			baseFound = true
			baseAmount = decimal.NewFromUint64(amount)
		} else if to.Equals(token1User) && from.Equals(token1Vault) && !quoteFound {
			quoteFound = true
			quoteAmount = decimal.NewFromUint64(amount)
		}
		if baseFound && quoteFound {
			break
		}
	}
	if !baseFound && !quoteFound {
		return nil, increaseOffset(offset), InstructionIgnoredError
	}

	baseTokenBalance, err := getTokenBalanceAfterTx(tx.rawTx, instruction.Accounts[4])
	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[5])
	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)

	event := "remove_liquidity"
	if !baseFound || !quoteFound {
		offset[1] += 1
		event = "remove_liquidity_one_side"
	} else {
		offset[1] += 2
	}
	return []Swap{
		{
			Program:           SolProgramRaydiumCPMM,
			Event:             event,
			Pool:              market,
			BaseMint:          token0,
			QuoteMint:         token1,
			BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
			QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
			BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
			QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
			BaseReserve:       baseReserve,
			QuoteReserve:      quoteReserve,
			BaseAmount:        baseAmount,
			QuoteAmount:       quoteAmount,
			User:              tx.rawTx.accountList[instruction.Accounts[0]],
			EntryContract:     entryContract,
		},
	}, offset, nil
}

func raydiumCPmmSwapParser(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 SwapBaseInputParser instruction")
	}

	var entryContract = tx.rawTx.accountList[tx.rawTx.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	discriminator := *(*[8]byte)(instruction.Data[:8])
	var swapMode SwapMode
	var fixedAmount decimal.Decimal
	var limitAmount decimal.Decimal
	switch discriminator {
	case raydiumCPmmSwapBaseInputDiscriminator:
		var args raydiumCPmmSwapBaseInputArgs
		if err := agbinary.NewBorshDecoder(instruction.Data[8:]).Decode(&args); err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("failed to decode raydium cpmm swap_base_input args: %w", err)
		}
		swapMode = SwapModeExactIn
		fixedAmount = decimal.NewFromUint64(args.AmountIn)
		limitAmount = decimal.NewFromUint64(args.MinimumAmountOut)
	case raydiumCPmmSwapBaseOutputDiscriminator:
		var args raydiumCPmmSwapBaseOutputArgs
		if err := agbinary.NewBorshDecoder(instruction.Data[8:]).Decode(&args); err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("failed to decode raydium cpmm swap_base_output args: %w", err)
		}
		swapMode = SwapModeExactOut
		fixedAmount = decimal.NewFromUint64(args.AmountOut)
		limitAmount = decimal.NewFromUint64(args.MaxAmountIn)
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}
	market := tx.rawTx.accountList[instruction.Accounts[3]]
	// Get token accounts from instruction
	tokenIn := tx.rawTx.accountList[instruction.Accounts[4]]
	tokenOut := tx.rawTx.accountList[instruction.Accounts[5]]
	user := tx.rawTx.accountList[instruction.Accounts[0]]

	user0 := instruction.Accounts[4]
	user1 := instruction.Accounts[5]

	inputVault := tx.rawTx.accountList[instruction.Accounts[6]]
	outputVault := tx.rawTx.accountList[instruction.Accounts[7]]

	vault0 := instruction.Accounts[6]
	vault1 := instruction.Accounts[7]

	inputTokenMint := tx.rawTx.accountList[instruction.Accounts[10]]
	outputTokenMint := tx.rawTx.accountList[instruction.Accounts[11]]

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

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

	baseReserve, _ := decimal.NewFromString(baseTokenBalance.UITokenAmount.Amount)
	quoteReserve, _ := decimal.NewFromString(quoteTokenBalance.UITokenAmount.Amount)
	userBase := getAccountBalanceAfterTx(tx.rawTx, user0)
	userQuote := getAccountBalanceAfterTx(tx.rawTx, user1)

	inners, err := getInnerInstructions(innerInstructions, offset[1])

	var baseFound, quoteFound bool
	var baseAmount, quoteAmount decimal.Decimal
	for _, inner := range inners {
		from, to, amount, err := parseTokenTransfer(tx.rawTx, inner)
		if err != nil {
			continue
		}
		if from.Equals(tokenIn) && to.Equals(inputVault) && !baseFound {
			baseFound = true
			baseAmount = decimal.NewFromUint64(amount)
		} else if from.Equals(outputVault) && to.Equals(tokenOut) && !quoteFound {
			quoteFound = true
			quoteAmount = decimal.NewFromUint64(amount)
		}
		if baseFound && quoteFound {
			break
		}
	}
	if !baseFound || !quoteFound {
		return nil, increaseOffset(offset), fmt.Errorf("failed to find token transfer in inner instructions")
	}
	offset[1] += 2
	swap := Swap{
		Program:           SolProgramRaydiumCPMM,
		Event:             "sell",
		Pool:              market,
		BaseMint:          inputTokenMint,
		QuoteMint:         outputTokenMint,
		BaseTokenProgram:  baseTokenBalance.ProgramIDAccount,
		QuoteTokenProgram: quoteTokenBalance.ProgramIDAccount,
		BaseMintDecimals:  uint8(baseTokenBalance.UITokenAmount.Decimals),
		QuoteMintDecimals: uint8(quoteTokenBalance.UITokenAmount.Decimals),
		BaseReserve:       baseReserve,
		QuoteReserve:      quoteReserve,
		BaseAmount:        baseAmount,
		QuoteAmount:       quoteAmount,
		User:              user,
		UserBaseBalance:   userBase,
		UserQuoteBalance:  userQuote,
		EntryContract:     entryContract,
	}
	swap.SetSwapAmountInfo(swapMode, fixedAmount, limitAmount)

	return []Swap{swap}, offset, nil
}