package shreder

import (
	"encoding/binary"
	"fmt"

	"github.com/gagliardetto/solana-go"
	"github.com/near/borsh-go"
)

var pumpAmmProgramID = solana.MustPublicKeyFromBase58("pAMMBay6oceH9fJKBRHGP5D4bD4sWpmSwMn52FMfXEA")
var (
	pumpAmmBuyTokensV2IX = []byte{198, 46, 21, 82, 180, 217, 232, 112}
	pumpAmmBuyTokensIX   = []byte{102, 6, 61, 18, 1, 218, 235, 234}
	pumpAmmSellTokensIX  = []byte{51, 230, 133, 164, 1, 127, 131, 173}
)

type pumpAmmBuyArgs struct {
	Amount     uint64
	MaxSolCost uint64
}

func parsePumpAmmInstruction(tx VersionedTransaction, instructionIndex int) (TxSignalBatch, error) {
	if instructionIndex >= len(tx.Instructions) {
		return nil, fmt.Errorf("instruction index out of bounds")
	}

	instruction := tx.Instructions[instructionIndex]
	if len(instruction.Data) == 0 {
		return nil, fmt.Errorf("data is empty")
	}

	var (
		err      error
		txSignal *TxSignal
	)
	if matchMethod(instruction.Data, pumpAmmBuyTokensIX) || matchMethod(instruction.Data, pumpAmmBuyTokensV2IX) {
		txSignal, err = parsePumpAmmBuy(tx, instruction)
	} else if matchMethod(instruction.Data, pumpAmmSellTokensIX) {
		txSignal, err = parsePumpAmmSell(tx, instruction)
	}

	if txSignal != nil {
		return TxSignalBatch{txSignal}, err
	}

	return nil, err

}

func decodePumpAmmBuyArgs(data []byte) (uint64, uint64, error) {
	if len(data) < 9 {
		return 0, 0, fmt.Errorf("data too short for pump amm buy args, len=%d", len(data))
	}

	var args pumpAmmBuyArgs
	if err := borsh.Deserialize(&args, data[8:]); err == nil {
		return args.Amount, args.MaxSolCost, nil
	}

	if len(data) >= 24 {
		amount := binary.LittleEndian.Uint64(data[8:16])
		maxSol := binary.LittleEndian.Uint64(data[16:24])
		return amount, maxSol, nil
	}

	return 0, 0, fmt.Errorf("failed to parse buy tokens args")
}

func parsePumpAmmBuy(tx VersionedTransaction, instruction Instructions) (*TxSignal, error) {
	amount, maxSol, err := decodePumpAmmBuyArgs(instruction.Data)
	if err != nil {
		return nil, err
	}
	exactIn := false
	if matchMethod(instruction.Data, pumpAmmBuyTokensV2IX) {
		temp := amount
		amount = maxSol
		maxSol = temp
		exactIn = true
	}

	if len(instruction.Accounts) < 7 {
		return nil, fmt.Errorf("accounts too short")
	}

	base, err := tx.GetAccount(int(instruction.Accounts[3]))
	if err != nil {
		return nil, err
	}
	quote, err := tx.GetAccount(int(instruction.Accounts[4]))
	if err != nil {
		return nil, err
	}
	if !quote.Equals(solana.WrappedSol) {
		return nil, nil
	}

	buyer, err := tx.GetAccount(int(instruction.Accounts[1]))
	if err != nil {
		return nil, err
	}

	return &TxSignal{
		TxHash:             tx.Signatures[0].String(),
		Label:              "pumpamm",
		Maker:              buyer.String(),
		Token0Address:      base.String(),
		Token1Address:      wsolMint,
		Token0Amount:       formatTokenAmount(amount),
		Token1Amount:       formatSolAmount(maxSol),
		Program:            "PumpAMM",
		Event:              "buy",
		IsToken2022:        false,
		IsMayhemMode:       false,
		ExactSOL:           exactIn,
		Block:              tx.Block,
		Token0AmountUint64: amount,
		Token1AmountUint64: maxSol,
	}, nil
}

func parsePumpAmmSell(tx VersionedTransaction, instruction Instructions) (*TxSignal, error) {
	amount, minSol, err := decodePumpAmmBuyArgs(instruction.Data)
	if err != nil {
		return nil, err
	}

	if len(instruction.Accounts) < 7 {
		return nil, fmt.Errorf("accounts too short")
	}

	base, err := tx.GetAccount(int(instruction.Accounts[3]))
	if err != nil {
		return nil, err
	}
	quote, err := tx.GetAccount(int(instruction.Accounts[4]))
	if err != nil {
		return nil, err
	}
	if !quote.Equals(solana.WrappedSol) {
		return nil, nil
	}

	buyer, err := tx.GetAccount(int(instruction.Accounts[1]))
	if err != nil {
		return nil, err
	}

	return &TxSignal{
		TxHash:             tx.Signatures[0].String(),
		Label:              "pumpamm",
		Maker:              buyer.String(),
		Token0Address:      base.String(),
		Token1Address:      wsolMint,
		Token0Amount:       formatTokenAmount(amount),
		Token1Amount:       formatSolAmount(minSol),
		Program:            "PumpAMM",
		Event:              "sell",
		IsToken2022:        false,
		IsMayhemMode:       false,
		Block:              tx.Block,
		Token0AmountUint64: amount,
		Token1AmountUint64: minSol,
	}, nil
}