package pump_parser

import (
	"bytes"
	"fmt"

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

func increaseOffset(offset [2]uint) [2]uint {
	if offset[1] == 0 {
		return [2]uint{offset[0] + 1, offset[1]}
	}
	return [2]uint{offset[0], offset[1] + 1}
}

// pumpParser // routes pump program instructions to their respective parsers,
// offset is [outerIndex, innerIndex] index of instructions in the transaction,
// if innerIndex == 0 this is outer instruction,if it's an inner instruction, outerIndex is the index of the parent instruction.
func pumpParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {

	if !tx.rawTx.accountList[instruction.ProgramIDIndex].Equals(pumpProgram) {
		return nil, increaseOffset(offset), fmt.Errorf("pump program instruction not found,  offset, %d, %d", offset[0], offset[1])
	}

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

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

	switch discriminator {
	case pumpBuyV2Discriminator, pumpBuyDiscriminator, pumpSellDiscriminator:
		return BuyOrSellParser(tx, instruction, innerInstructions, offset)
	case pumpCreateDiscriminator, pumpCreateV2Discriminator:
		return CreateParser(tx, instruction, innerInstructions, offset)
	case pumpMigrateDiscriminator:
		return MigrateParser(tx, instruction, innerInstructions, offset)
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}
}

type PumpCreateData struct {
	Discriminator uint64
	Name          string
	Symbol        string
	Uri           string
	Creator       solana.PublicKey
}

type PumpCreateV2Data struct {
	Discriminator uint64
	Name          string
	Symbol        string
	Uri           string

	Creator  solana.PublicKey
	IsMayhem bool
}

type PumpCreateEvent struct {
	Name   string
	Symbol string
	Uri    string

	Mint         solana.PublicKey
	BondingCurve solana.PublicKey
	User         solana.PublicKey
	Creator      solana.PublicKey

	Timestamp            int64
	VirtualTokenReserves uint64
	VirtualSolReserves   uint64
	RealTokenReserves    uint64
	TokenTotalSupply     uint64
	TokenProgram         solana.PublicKey
	IsMayhemMode         bool
}

func CreateParser(tx *Tx, instr Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	result := tx.rawTx
	var entryContract = result.accountList[result.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]

	var programIndex = instr.ProgramIDIndex
	var err error

	var createEvent PumpCreateEvent
	var prefixLen = offset[1]
	inners, err := getInnerInstructions(innerInstructions, prefixLen)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("pump create get inner instructions error: %v, offset, %d, %d", err, offset[0], offset[1])
	}
	for innerIndex, innerInstr := range inners {
		if innerInstr.ProgramIDIndex == programIndex && bytes.Equal(innerInstr.Data[:8], pumpEventDiscriminator[:]) && bytes.Equal(innerInstr.Data[8:16], pumpCreateEventDiscriminator[:]) {
			err = agbinary.NewBorshDecoder(innerInstr.Data[16:]).Decode(&createEvent)
			if offset[1] == 0 {
				offset[0] += 1
			} else {
				offset[1] = uint(innerIndex) + 1 + prefixLen
			}
			if err != nil {
				return nil, offset, fmt.Errorf("pump create event decode error: %v, offset, %d, %d", err, offset[0], offset[1])
			}
			break
		}
	}
	if createEvent == (PumpCreateEvent{}) {
		return nil, increaseOffset(offset), fmt.Errorf("pump create event not found, offset, %d, %d", offset[0], offset[1])
	}
	userIndex := 0
	if bytes.HasPrefix(instr.Data, pumpCreateV2Discriminator[:]) {
		userIndex = instr.Accounts[5]
	} else if bytes.HasPrefix(instr.Data, pumpCreateDiscriminator[:]) {
		userIndex = instr.Accounts[7]
	}
	userBase := getAccountBalanceAfterTx(result, userIndex)
	userQuote, _ := GetSolAfterTx(result, userIndex)

	totalSupply := decimal.NewFromUint64(createEvent.TokenTotalSupply).Div(decimal.New(1, 6))
	tx.Token[createEvent.Mint] = TokenMeta{
		Mint:         createEvent.Mint,
		TokenProgram: createEvent.TokenProgram,
		Decimals:     6,
		Name:         createEvent.Name,
		Symbol:       createEvent.Symbol,
		Url:          createEvent.Uri,
		TotalSupply:  &totalSupply,
	}
	return []Swap{
		{
			Program:           SolProgramPump,
			Event:             "create",
			Pool:              createEvent.BondingCurve,
			BaseMint:          createEvent.Mint,
			QuoteMint:         solana.PublicKey{},
			BaseTokenProgram:  createEvent.TokenProgram,
			QuoteTokenProgram: solana.PublicKey{},
			Creator:           createEvent.Creator,
			BaseMintDecimals:  6,
			QuoteMintDecimals: 9,
			User:              createEvent.User,
			BaseAmount:        decimal.Zero,
			QuoteAmount:       decimal.Zero,
			BaseReserve:       decimal.NewFromUint64(createEvent.RealTokenReserves),
			QuoteReserve:      decimal.Zero,
			Mayhem:            createEvent.IsMayhemMode,
			UserBaseBalance:   userBase,
			UserQuoteBalance:  decimal.NewFromUint64(userQuote),
			EntryContract:     entryContract,
		},
	}, offset, nil
}

type PumpTradeEvent struct {
	Mint                 solana.PublicKey
	SolAmount            uint64
	TokenAmount          uint64
	IsBuy                bool
	User                 solana.PublicKey
	Timestamp            int64
	VirtualSolReserves   uint64
	VirtualTokenReserves uint64

	RealSolReserves   uint64
	RealTokenReserves uint64

	FeeRecipient   solana.PublicKey
	FeeBasisPoints uint64
	Fee            uint64

	Creator solana.PublicKey

	CreatorFeeBasisPoints uint64
	CreatorFee            uint64
}

type PumpTradeFeeArg struct {
	IsPump    bool
	MarketCap [16]byte
	TradeSize uint64
}

type CompleteEvent struct {
	User         solana.PublicKey
	Mint         solana.PublicKey
	BondingCurve solana.PublicKey
	Timestamp    int64
}

func BuyOrSellParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	result := tx.rawTx
	var entryContract = result.accountList[result.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	var err error
	var programIndex = instruction.ProgramIDIndex

	feeEventProgramIndex := 0
	for i, b := range result.accountList {
		if b.Equals(pumpFeesProgram) {
			feeEventProgramIndex = i
			break
		}
	}
	var (
		tradeEvent    PumpTradeEvent
		tradeFeeArg   PumpTradeFeeArg
		completeEvent CompleteEvent
		completed     bool
		newoffset     [2]uint
	)

	var prefixLen = offset[1]
	inners, err := getInnerInstructions(innerInstructions, prefixLen)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("pump create get inner instructions error: %v,offset, %d, %d", err, offset[0], offset[1])
	}
	if instruction.StackHeight != nil && *instruction.StackHeight > 2 {
		for _, innerInstr := range innerInstructions.Instructions {
			if innerInstr.StackHeight != nil && *innerInstr.StackHeight == *instruction.StackHeight-1 {
				entryContract = result.accountList[innerInstr.ProgramIDIndex]
			}
		}
	}

	for innerIndex, innerInstr := range inners {
		if innerInstr.ProgramIDIndex == feeEventProgramIndex && bytes.Equal(innerInstr.Data[:8], pumpGetFeesDiscriminator[:]) {
			err = agbinary.NewBorshDecoder(innerInstr.Data[8:]).Decode(&tradeFeeArg)
			if err != nil {
				return nil, increaseOffset(offset), fmt.Errorf("pump get fees event decode error: %v, offset, %d, %d", err, offset[0], offset[1])
			}
			continue
		}
		if innerInstr.ProgramIDIndex == programIndex && bytes.Equal(innerInstr.Data[:8], pumpEventDiscriminator[:]) {
			if bytes.Equal(innerInstr.Data[8:16], pumpTradeEventDiscriminator[8:16]) {
				err = agbinary.NewBorshDecoder(innerInstr.Data[16:]).Decode(&tradeEvent)
				if offset[1] == 0 {
					newoffset = [2]uint{offset[0] + 1, offset[1]}
				} else {
					newoffset = [2]uint{offset[0], prefixLen + uint(innerIndex) + 1}
				}
				if err != nil {
					return nil, newoffset, fmt.Errorf("pump buy event decode error: %v, offset, %d, %d", err, offset[0], offset[1])
				}
				if !tradeEvent.IsBuy {
					break
				}
			} else if bytes.Equal(innerInstr.Data[8:16], pumpCompleteEventDiscriminator[:]) {
				err = agbinary.NewBorshDecoder(innerInstr.Data[16:]).Decode(&completeEvent)
				if offset[1] == 0 {
					newoffset = [2]uint{offset[0] + 1, offset[1]}
				} else {
					newoffset = [2]uint{offset[0], prefixLen + uint(innerIndex) + 1}
				}
				if err != nil {
					return nil, newoffset, fmt.Errorf("pump completeEvent event decode error: %v, offset, %d, %d", err, offset[0], offset[1])
				}
				completed = true
				break
			}

		}
	}
	if tradeEvent == (PumpTradeEvent{}) {
		return nil, increaseOffset(offset), fmt.Errorf("pmp buy/sell event not found, offset, %d, %d", offset[0], offset[1])
	}

	offset = [2]uint{newoffset[0], newoffset[1]}

	event := ""
	baseTokenProgram := solana.TokenProgramID
	if tradeEvent.IsBuy {
		event = "buy"
		baseTokenProgram = result.accountList[instruction.Accounts[8]]
	} else {
		event = "sell"
		baseTokenProgram = result.accountList[instruction.Accounts[9]]
	}
	if _, exists := tx.Token[tradeEvent.Mint]; !exists {
		tx.Token[tradeEvent.Mint] = TokenMeta{
			Mint:         tradeEvent.Mint,
			TokenProgram: baseTokenProgram,
			Decimals:     6,
		}
	}

	var user = tradeEvent.User

	ataUserIdx := instruction.Accounts[5]
	userIndex := instruction.Accounts[6]
	if !tradeEvent.User.IsOnCurve() && (entryContract.Equals(okxDexRoutersV2) || entryContract.Equals(okxAggregatorV2)) {
		userBaseAmount, ataIndex := tokenBalanceChange(result, 0, baseTokenProgram, tradeEvent.Mint)
		//&& userBaseAmount.BigInt().Uint64() == tradeEvent.TokenAmount
		if !userBaseAmount.IsZero() {
			user = result.accountList[0]
			userIndex = 0
			ataUserIdx = ataIndex
		}
	}

	userBase := getAccountBalanceAfterTx(result, ataUserIdx)
	userQuote, _ := GetSolAfterTx(result, userIndex)

	solAmount := tradeEvent.SolAmount
	if tradeEvent.IsBuy && bytes.Equal(instruction.Data[:8], pumpBuyV2Discriminator[:]) {
		fee := tradeEvent.Fee + tradeEvent.CreatorFee
		solAmount = tradeFeeArg.TradeSize
		if solAmount > fee {
			solAmount = solAmount - fee
		}
	}
	swaps := []Swap{
		{
			Program:           SolProgramPump,
			Event:             event,
			Pool:              result.accountList[instruction.Accounts[3]],
			BaseMint:          tradeEvent.Mint,
			QuoteMint:         solana.PublicKey{},
			BaseTokenProgram:  baseTokenProgram,
			QuoteTokenProgram: solana.PublicKey{},
			Creator:           tradeEvent.Creator,
			BaseMintDecimals:  6,
			QuoteMintDecimals: 9,
			User:              user,
			BaseAmount:        decimal.NewFromUint64(tradeEvent.TokenAmount),
			QuoteAmount:       decimal.NewFromUint64(solAmount),
			BaseReserve:       decimal.NewFromUint64(tradeEvent.RealTokenReserves),
			QuoteReserve:      decimal.NewFromUint64(tradeEvent.RealSolReserves),
			Mayhem:            isMayhemPump(result.accountList[instruction.Accounts[1]]),
			UserBaseBalance:   userBase,
			UserQuoteBalance:  decimal.NewFromUint64(userQuote),
			EntryContract:     entryContract,
		},
	}
	if completed {
		swaps = append(swaps, Swap{
			Program:           SolProgramPump,
			Event:             "complete",
			Pool:              result.accountList[instruction.Accounts[3]],
			BaseMint:          tradeEvent.Mint,
			QuoteMint:         solana.PublicKey{},
			BaseTokenProgram:  result.accountList[instruction.Accounts[8]],
			QuoteTokenProgram: solana.PublicKey{},
			Creator:           tradeEvent.Creator,
			BaseMintDecimals:  6,
			QuoteMintDecimals: 9,
			User:              user,
			Mayhem:            isMayhemPump(result.accountList[instruction.Accounts[1]]),
			UserBaseBalance:   userBase,
			UserQuoteBalance:  decimal.NewFromUint64(userQuote),
			EntryContract:     entryContract,
		})
	}

	return swaps, offset, nil
}

type MigrateEvent struct {
	User solana.PublicKey
	Mint solana.PublicKey
	//Creator          solana.PublicKey

	MintAmount       uint64
	SolAmount        uint64
	PoolMigrationFee uint64
	//CreatorFee       uint64
	BondingCurve solana.PublicKey
	TimeStamp    int64
	Pool         solana.PublicKey
}

func MigrateParser(tx *Tx, instr Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	result := tx.rawTx
	var entryContract = result.accountList[result.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	var err error
	programIndex := instr.ProgramIDIndex
	ammprogramIdx := 0
	for i, b := range result.accountList {
		if b.Equals(pumpAmmProgram) {
			ammprogramIdx = i
			break
		}
	}
	if ammprogramIdx == 0 {
		return nil, increaseOffset(offset), fmt.Errorf("pump migrate, amm program id not found in account list, offset, %d, %d", offset[0], offset[1])
	}
	var prefixLen = offset[1]
	inners, err := getInnerInstructions(innerInstructions, prefixLen)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("pump migrate get inner instructions  offset, %d, %d", offset[0], offset[1])
	}
	var (
		migrateEvent MigrateEvent
		createEvent  ammCreatePoolEvent
		newoffset    [2]uint
	)
	for innerIndex, innerInstr := range inners {
		if innerInstr.ProgramIDIndex == ammprogramIdx && bytes.Equal(innerInstr.Data[:8], pumpAmmEventDiscriminator[:]) {
			if bytes.Equal(innerInstr.Data[8:16], pumpAmmCreateEventDiscriminator[:]) {
				err = agbinary.NewBorshDecoder(innerInstr.Data[16:]).Decode(&createEvent)
				if offset[1] == 0 {
					newoffset = [2]uint{offset[0] + 1, offset[1]}
				} else {
					newoffset = [2]uint{offset[0], prefixLen + uint(innerIndex) + 1}
				}
				if err != nil {
					return nil, newoffset, fmt.Errorf("pump amm createEvent decode error: %v, offset, %d, %d", err, offset[0], offset[1])
				}
			}
			//
		} else if innerInstr.ProgramIDIndex == programIndex && bytes.Equal(innerInstr.Data[:8], pumpEventDiscriminator[:]) {
			if bytes.Equal(innerInstr.Data[8:16], pumpMigrateEventDiscriminator[:]) {
				err = agbinary.NewBorshDecoder(innerInstr.Data[16:]).Decode(&migrateEvent)
				if offset[1] == 0 {
					newoffset = [2]uint{offset[0] + 1, offset[1]}
				} else {
					newoffset = [2]uint{offset[0], prefixLen + uint(innerIndex) + 1}
				}
				if err != nil {
					return nil, newoffset, fmt.Errorf("pump buy event decode error: %v, offset, %d, %d", err, offset[0], offset[1])
				}
				break
			}
		}
	}
	if migrateEvent == (MigrateEvent{}) || createEvent == (ammCreatePoolEvent{}) {
		offset = [2]uint{newoffset[0], newoffset[1]}
		return nil, increaseOffset(offset), InstructionIgnoredError
	}

	offset = [2]uint{newoffset[0], newoffset[1]}
	// verify migrate by checking create pool and migrate event
	userIndex := instr.Accounts[5]
	ataBondingCurveAccountIndex := instr.Accounts[4]
	bc, err := getTokenBalanceAfterTx(result, ataBondingCurveAccountIndex)
	if err != nil || bc == nil {
		return nil, increaseOffset(offset), fmt.Errorf("pump migrate get bonding curve balance error: %v, offset, %d, %d", err, offset[0], offset[1])
	}
	baseTokenProgram := bc.ProgramIDAccount
	var userBase decimal.Decimal
	if result.accountList[userIndex].Equals(pumpMigrationAccount) {
		userBase = decimal.Zero
	} else {
		userBase = GetTokenBalanceAfterTx(result, userIndex, baseTokenProgram, migrateEvent.Mint)
	}
	userQuote, _ := GetSolAfterTx(result, userIndex)

	if _, exists := tx.Token[migrateEvent.Mint]; !exists {
		tx.Token[migrateEvent.Mint] = TokenMeta{
			Mint:         migrateEvent.Mint,
			TokenProgram: baseTokenProgram,
			Decimals:     6,
		}
	}
	swaps := []Swap{
		{
			Program:           SolProgramPump,
			Event:             "migrate",
			Pool:              migrateEvent.BondingCurve,
			BaseMint:          migrateEvent.Mint,
			QuoteMint:         solana.PublicKey{},
			BaseTokenProgram:  baseTokenProgram,
			QuoteTokenProgram: solana.PublicKey{},
			Creator:           createEvent.Creator,
			BaseMintDecimals:  6,
			QuoteMintDecimals: 9,
			User:              migrateEvent.User,
			//BaseAmount:        decimal.Decimal{},
			//QuoteAmount:       decimal.Decimal{},
			BaseReserve:       decimal.NewFromUint64(migrateEvent.MintAmount),
			QuoteReserve:      decimal.NewFromUint64(migrateEvent.SolAmount),
			Mayhem:            createEvent.IsMayhemMode,
			MigrateTopProgram: pumpAmmProgram,
			MigrateToPool:     migrateEvent.Pool,
			UserBaseBalance:   userBase,
			UserQuoteBalance:  decimal.NewFromUint64(userQuote),
			EntryContract:     entryContract,
		},
	}
	swaps = append(swaps, Swap{
		Program:           SolProgramPumpAMM,
		Event:             "create",
		Pool:              migrateEvent.Pool,
		BaseMint:          migrateEvent.Mint,
		QuoteMint:         wSolMint,
		BaseTokenProgram:  baseTokenProgram,
		QuoteTokenProgram: solana.TokenProgramID,
		Creator:           createEvent.Creator,
		BaseMintDecimals:  6,
		QuoteMintDecimals: 9,
		User:              migrateEvent.User,
		BaseAmount:        decimal.NewFromUint64(migrateEvent.MintAmount),
		QuoteAmount:       decimal.NewFromUint64(migrateEvent.SolAmount),
		BaseReserve:       decimal.NewFromUint64(migrateEvent.MintAmount),
		QuoteReserve:      decimal.NewFromUint64(migrateEvent.SolAmount),
		Mayhem:            createEvent.IsMayhemMode,
		UserBaseBalance:   userBase,
		UserQuoteBalance:  decimal.NewFromUint64(userQuote),
		EntryContract:     entryContract,
	})

	return swaps, offset, nil

}