package pump_parser

import (
	"bytes"
	"fmt"

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

type meteoraDlmmSwapArgs struct {
	AmountIn     uint64
	MinAmountOut uint64
}

type meteoraDlmmSwapExactOutArgs struct {
	MaxInAmount uint64
	OutAmount   uint64
}

type meteoraDlmmSwapWithPriceImpactArgs struct {
	AmountIn          uint64
	ActiveID          *int32 `bin:"optional"`
	MaxPriceImpactBps uint16
}

type dlmmSwapEvent struct {
	LbPair      solana.PublicKey
	From        solana.PublicKey
	StartBinId  int32
	EndBinId    int32
	AmountIn    uint64
	AmountOut   uint64
	SwapForY    bool
	Fee         uint64
	ProtocolFee uint64
	FeeBps      agbinary.Uint128
	HostFee     uint64
}

type dlmmSwapAccounts struct {
	poolIdx          int
	reserveXIdx      int
	reserveYIdx      int
	userTokenInIdx   int
	userTokenOutIdx  int
	tokenXMintIdx    int
	tokenYMintIdx    int
	oracleIdx        int
	userIdx          int
	tokenXProgramIdx int
	tokenYProgramIdx int
}

var meteoraDlmmEventAuthority = func() solana.PublicKey {
	key, _, err := solana.FindProgramAddress([][]byte{[]byte("__event_authority")}, meteoraDlmmProgram)
	if err != nil {
		return solana.PublicKey{}
	}
	return key
}()

func metaoradlmmParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	if !tx.rawTx.accountList[instruction.ProgramIDIndex].Equals(meteoraDlmmProgram) {
		return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm program instruction not found, offset, %d, %d", offset[0], offset[1])
	}
	decode := instruction.Data
	if len(decode) < 8 {
		offset[1] += 1
		return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm program instruction data too short, offset, %d, %d", offset[0], offset[1])
	}

	discriminator := *(*[8]byte)(decode[:8])
	switch discriminator {
	case meteoraDlmmSwapDiscriminator, meteoraDlmmSwapExactOutDiscriminator, meteoraDlmmSwapWithPriceImpactDiscriminator:
		return metaoradlmmSwapParser(tx, instruction, innerInstructions, offset)
	case meteoraDlmmSwap2Discriminator, meteoraDlmmSwapExactOut2Discriminator, meteoraDlmmSwapWithPriceImpact2Discriminator:
		return metaoradlmmSwap2Parser(tx, instruction, innerInstructions, offset)
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}
}

func metaoradlmmSwapParser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	result := tx.rawTx

	entryContract := result.accountList[result.Transaction.Message.Instructions[offset[0]].ProgramIDIndex]
	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]
			}
		}
	}

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

	discriminator := *(*[8]byte)(decode[:8])
	switch discriminator {
	case meteoraDlmmSwapDiscriminator, meteoraDlmmSwap2Discriminator:
		var args meteoraDlmmSwapArgs
		if err := agbinary.NewBorshDecoder(decode[8:]).Decode(&args); err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm swap decode error: %v, offset, %d, %d", err, offset[0], offset[1])
		}
	case meteoraDlmmSwapExactOutDiscriminator, meteoraDlmmSwapExactOut2Discriminator:
		var args meteoraDlmmSwapExactOutArgs
		if err := agbinary.NewBorshDecoder(decode[8:]).Decode(&args); err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm swap_exact_out decode error: %v, offset, %d, %d", err, offset[0], offset[1])
		}
	case meteoraDlmmSwapWithPriceImpactDiscriminator, meteoraDlmmSwapWithPriceImpact2Discriminator:
		var args meteoraDlmmSwapWithPriceImpactArgs
		if err := agbinary.NewBorshDecoder(decode[8:]).Decode(&args); err != nil {
			return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm swap_with_price_impact decode error: %v, offset, %d, %d", err, offset[0], offset[1])
		}
	default:
		return nil, increaseOffset(offset), InstructionIgnoredError
	}

	accounts, err := resolveDlmmSwapAccounts(result, instruction.Accounts)
	if err != nil {
		return nil, increaseOffset(offset), fmt.Errorf("meteora dlmm swap accounts parse error: %v, offset, %d, %d", err, offset[0], offset[1])
	}
	pool := result.accountList[accounts.poolIdx]
	reserveXIdx := accounts.reserveXIdx
	reserveYIdx := accounts.reserveYIdx
	userTokenInIdx := accounts.userTokenInIdx
	userTokenOutIdx := accounts.userTokenOutIdx
	tokenXMint := result.accountList[accounts.tokenXMintIdx]
	tokenYMint := result.accountList[accounts.tokenYMintIdx]
	userIdx := accounts.userIdx
	tokenXProgram := result.accountList[accounts.tokenXProgramIdx]
	tokenYProgram := result.accountList[accounts.tokenYProgramIdx]

	swapEvent, nextOffset, err := dlmmSwapEventFromInnerInstructions(innerInstructions, instruction, offset)
	if err != nil {
		return nil, nextOffset, err
	}
	offset = nextOffset

	baseMint, quoteMint, baseIsX := dlmmSelectBaseQuote(tokenXMint, tokenYMint)
	baseTokenProgram := tokenXProgram
	quoteTokenProgram := tokenYProgram
	baseReserveIdx := reserveXIdx
	quoteReserveIdx := reserveYIdx
	if !baseIsX {
		baseTokenProgram = tokenYProgram
		quoteTokenProgram = tokenXProgram
		baseReserveIdx = reserveYIdx
		quoteReserveIdx = reserveXIdx
	}

	swapForY := swapEvent.SwapForY
	inputIsX := swapForY
	amountIn := decimal.NewFromUint64(swapEvent.AmountIn)
	amountOut := decimal.NewFromUint64(swapEvent.AmountOut)

	event := "buy"
	if baseIsX == inputIsX {
		event = "sell"
	}

	userBaseIdx := userTokenOutIdx
	userQuoteIdx := userTokenInIdx
	if baseIsX == inputIsX {
		userBaseIdx = userTokenInIdx
		userQuoteIdx = userTokenOutIdx
	}

	baseAmount := amountOut
	quoteAmount := amountIn
	if baseIsX {
		if swapForY {
			baseAmount = amountIn
			quoteAmount = amountOut
		}
	} else {
		if !swapForY {
			baseAmount = amountIn
			quoteAmount = amountOut
		}
	}

	eventUser := result.accountList[userIdx]
	if !swapEvent.From.IsZero() {
		eventUser = swapEvent.From
	}
	if !eventUser.IsOnCurve() && (entryContract.Equals(okxDexRoutersV2) || entryContract.Equals(okxAggregatorV2)) {
		userBaseAmount, ataIndex := tokenBalanceChange(result, 0, baseTokenProgram, baseMint)
		if !userBaseAmount.IsZero() {
			eventUser = result.accountList[0]
			userIdx = 0
			if ataIndex > 0 {
				userBaseIdx = ataIndex
			}
		}
	}

	baseDecimals, ok := dlmmTokenDecimals(result, baseReserveIdx)
	if !ok {
		baseDecimals, _ = dlmmTokenDecimals(result, userBaseIdx)
	}
	quoteDecimals, ok := dlmmTokenDecimals(result, quoteReserveIdx)
	if !ok {
		quoteDecimals, _ = dlmmTokenDecimals(result, userQuoteIdx)
	}

	if _, exists := tx.Token[baseMint]; !exists && !baseMint.Equals(wSolMint) {
		tx.Token[baseMint] = TokenMeta{
			Mint:         baseMint,
			Decimals:     baseDecimals,
			TokenProgram: baseTokenProgram,
		}
	}
	if _, exists := tx.Token[quoteMint]; !exists && !quoteMint.Equals(wSolMint) {
		tx.Token[quoteMint] = TokenMeta{
			Mint:         quoteMint,
			Decimals:     quoteDecimals,
			TokenProgram: quoteTokenProgram,
		}
	}

	baseReserve := getAccountBalanceAfterTx(result, baseReserveIdx)
	quoteReserve := getAccountBalanceAfterTx(result, quoteReserveIdx)
	userBase := getAccountBalanceAfterTx(result, userBaseIdx)
	userQuote := GetTokenBalanceAfterTx(result, userIdx, quoteTokenProgram, quoteMint)
	if quoteMint.Equals(wSolMint) {
		if solAmount, err := GetSolAfterTx(result, userIdx); err == nil {
			userQuote = userQuote.Add(decimal.NewFromUint64(solAmount))
		}
	}

	swap := Swap{
		Program:           SolProgramMeteoraDLMM,
		Event:             event,
		Pool:              pool,
		BaseMint:          baseMint,
		QuoteMint:         quoteMint,
		BaseTokenProgram:  baseTokenProgram,
		QuoteTokenProgram: quoteTokenProgram,
		Creator:           solana.PublicKey{},
		BaseMintDecimals:  baseDecimals,
		QuoteMintDecimals: quoteDecimals,
		User:              eventUser,
		BaseAmount:        baseAmount,
		QuoteAmount:       quoteAmount,
		BaseReserve:       baseReserve,
		QuoteReserve:      quoteReserve,
		UserBaseBalance:   userBase,
		UserQuoteBalance:  userQuote,
		EntryContract:     entryContract,
		StartBinId:        swapEvent.StartBinId,
		EndBinId:          swapEvent.EndBinId,
	}

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

func metaoradlmmSwap2Parser(tx *Tx, instruction Instruction, innerInstructions InnerInstructions, offset [2]uint) ([]Swap, [2]uint, error) {
	return metaoradlmmSwapParser(tx, instruction, innerInstructions, offset)
}

func dlmmSelectBaseQuote(tokenX, tokenY solana.PublicKey) (baseMint solana.PublicKey, quoteMint solana.PublicKey, baseIsX bool) {
	priority := []solana.PublicKey{wSolMint, usdcMint, usd1Mint}
	for _, mint := range priority {
		if tokenX.Equals(mint) {
			return tokenY, tokenX, false
		}
		if tokenY.Equals(mint) {
			return tokenX, tokenY, true
		}
	}
	return tokenX, tokenY, true
}

func dlmmSwapEventFromInnerInstructions(innerInstructions InnerInstructions, instruction Instruction, offset [2]uint) (dlmmSwapEvent, [2]uint, error) {
	var prefixLen = offset[1]
	inners, err := getInnerInstructions(innerInstructions, prefixLen)
	if err != nil {
		return dlmmSwapEvent{}, increaseOffset(offset), fmt.Errorf("meteora dlmm swap get inner instructions error: %v, offset, %d, %d", err, offset[0], prefixLen)
	}
	for innerIndex, innerInstr := range inners {
		if innerInstr.ProgramIDIndex != instruction.ProgramIDIndex {
			continue
		}
		event, ok := dlmmDecodeSwapEvent(innerInstr.Data)
		if !ok {
			continue
		}
		if offset[1] == 0 {
			offset[0] += 1
		} else {
			offset[1] += uint(innerIndex) + 1 + prefixLen
		}
		return event, offset, nil
	}
	return dlmmSwapEvent{}, increaseOffset(offset), fmt.Errorf("meteora dlmm swap event not found, offset, %d, %d", offset[0], prefixLen)
}

func dlmmDecodeSwapEvent(data []byte) (dlmmSwapEvent, bool) {
	switch {
	case len(data) >= 8 && bytes.Equal(data[:8], meteoraDlmmSwapEventDiscriminator[:]):
		var event dlmmSwapEvent
		if err := agbinary.NewBorshDecoder(data[8:]).Decode(&event); err != nil {
			return dlmmSwapEvent{}, false
		}
		return event, true
	case len(data) >= 16 &&
		bytes.Equal(data[:8], eventDiscriminator[:]) &&
		bytes.Equal(data[8:16], meteoraDlmmSwapEventDiscriminator[:]):
		var event dlmmSwapEvent
		if err := agbinary.NewBorshDecoder(data[16:]).Decode(&event); err != nil {
			return dlmmSwapEvent{}, false
		}
		return event, true
	default:
		return dlmmSwapEvent{}, false
	}
}

func resolveDlmmSwapAccounts(result *RawTx, accounts []int) (dlmmSwapAccounts, error) {
	if len(accounts) < 13 {
		return dlmmSwapAccounts{}, fmt.Errorf("accounts too short, expected at least 13")
	}
	accountList := result.accountList

	basePosCandidates := []int{1, 2}
	for _, basePos := range basePosCandidates {
		if basePos+6 >= len(accounts) {
			continue
		}
		oraclePos := basePos + 6

		userPos := oraclePos + 1
		hostFeePresent := true
		if userPos < len(accounts) && dlmmIsSigner(result, accounts[userPos]) {
			hostFeePresent = false
		} else {
			userPos = oraclePos + 2
		}
		if userPos+2 >= len(accounts) {
			continue
		}
		tokenXProgramPos := userPos + 1
		tokenYProgramPos := userPos + 2

		eventAuthorityPos := tokenYProgramPos + 1
		if eventAuthorityPos < len(accounts) && accountList[accounts[eventAuthorityPos]].Equals(solana.MemoProgramID) {
			eventAuthorityPos++
		}
		programPos := eventAuthorityPos + 1
		if programPos >= len(accounts) {
			continue
		}
		if !accountList[accounts[eventAuthorityPos]].Equals(meteoraDlmmEventAuthority) {
			continue
		}
		if !accountList[accounts[programPos]].Equals(meteoraDlmmProgram) {
			continue
		}

		if hostFeePresent && oraclePos+1 < len(accounts) && dlmmIsSigner(result, accounts[oraclePos+1]) {
			continue
		}

		return dlmmSwapAccounts{
			poolIdx:          accounts[0],
			reserveXIdx:      accounts[oraclePos-6],
			reserveYIdx:      accounts[oraclePos-5],
			userTokenInIdx:   accounts[oraclePos-4],
			userTokenOutIdx:  accounts[oraclePos-3],
			tokenXMintIdx:    accounts[oraclePos-2],
			tokenYMintIdx:    accounts[oraclePos-1],
			oracleIdx:        accounts[oraclePos],
			userIdx:          accounts[userPos],
			tokenXProgramIdx: accounts[tokenXProgramPos],
			tokenYProgramIdx: accounts[tokenYProgramPos],
		}, nil
	}

	return dlmmSwapAccounts{}, fmt.Errorf("accounts layout invalid")
}

func dlmmTokenDecimals(result *RawTx, accountIndex int) (uint8, bool) {
	for _, meta := range result.Meta.PostTokenBalances {
		if meta.AccountIndex == accountIndex {
			return uint8(meta.UITokenAmount.Decimals), true
		}
	}
	for _, meta := range result.Meta.PreTokenBalances {
		if meta.AccountIndex == accountIndex {
			return uint8(meta.UITokenAmount.Decimals), true
		}
	}
	return 0, false
}

func dlmmTokenDelta(result *RawTx, accountIndex int) (decimal.Decimal, bool) {
	before, okBefore := dlmmTokenAmount(result, accountIndex, false)
	after, okAfter := dlmmTokenAmount(result, accountIndex, true)
	if !okBefore && !okAfter {
		return decimal.Zero, false
	}
	if !okBefore {
		before = decimal.Zero
	}
	if !okAfter {
		after = decimal.Zero
	}
	return after.Sub(before).Abs(), true
}

func dlmmTokenDeltaSigned(result *RawTx, accountIndex int) (decimal.Decimal, bool) {
	before, okBefore := dlmmTokenAmount(result, accountIndex, false)
	after, okAfter := dlmmTokenAmount(result, accountIndex, true)
	if !okBefore && !okAfter {
		return decimal.Zero, false
	}
	if !okBefore {
		before = decimal.Zero
	}
	if !okAfter {
		after = decimal.Zero
	}
	return after.Sub(before), true
}

func dlmmTokenAmount(result *RawTx, accountIndex int, post bool) (decimal.Decimal, bool) {
	var balances []TokenBalance
	if post {
		balances = result.Meta.PostTokenBalances
	} else {
		balances = result.Meta.PreTokenBalances
	}
	for _, meta := range balances {
		if meta.AccountIndex == accountIndex {
			amount, err := decimal.NewFromString(meta.UITokenAmount.Amount)
			if err != nil {
				return decimal.Zero, false
			}
			return amount, true
		}
	}
	return decimal.Zero, false
}

func dlmmIsSigner(result *RawTx, accountIndex int) bool {
	if accountIndex < 0 || accountIndex >= len(result.Transaction.Message.AccountKeys) {
		return false
	}
	return accountIndex < result.Transaction.Message.Header.NumRequiredSignatures
}

func dlmmTokenBalanceMeta(result *RawTx, accountIndex int) (TokenBalance, bool) {
	for _, meta := range result.Meta.PostTokenBalances {
		if meta.AccountIndex == accountIndex {
			return meta, true
		}
	}
	for _, meta := range result.Meta.PreTokenBalances {
		if meta.AccountIndex == accountIndex {
			return meta, true
		}
	}
	return TokenBalance{}, false
}