package pump_parser

import (
	"bytes"
	"testing"

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

func testPublicKey(seed byte) solana.PublicKey {
	buf := make([]byte, solana.PublicKeyLength)
	for i := range buf {
		buf[i] = seed
	}
	return solana.PublicKeyFromBytes(buf)
}

func seqInts(n int) []int {
	out := make([]int, n)
	for i := range out {
		out[i] = i
	}
	return out
}

func mustBorshEncode(t *testing.T, value any) []byte {
	t.Helper()

	var buf bytes.Buffer
	if err := agbinary.NewBorshEncoder(&buf).Encode(value); err != nil {
		t.Fatalf("borsh encode failed: %v", err)
	}
	return buf.Bytes()
}

func TestMeteoraDlmmInitializeParserCompatibility(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name                string
		discriminator       [8]byte
		accountCount        int
		wantPoolPos         int
		wantBaseMintPos     int
		wantQuoteMintPos    int
		wantUserPos         int
		wantBaseProgramPos  int
		wantQuoteProgramPos int
	}{
		{
			name:                "initialize_lb_pair",
			discriminator:       meteoraInitializeLbPairDiscriminator,
			accountCount:        14,
			wantPoolPos:         0,
			wantBaseMintPos:     2,
			wantQuoteMintPos:    3,
			wantUserPos:         8,
			wantBaseProgramPos:  9,
			wantQuoteProgramPos: 9,
		},
		{
			name:                "initialize_lb_pair2",
			discriminator:       meteoraInitializeLbPair2Discriminator,
			accountCount:        16,
			wantPoolPos:         0,
			wantBaseMintPos:     2,
			wantQuoteMintPos:    3,
			wantUserPos:         8,
			wantBaseProgramPos:  11,
			wantQuoteProgramPos: 12,
		},
		{
			name:                "initialize_customizable_permissionless_lb_pair",
			discriminator:       meteoraInitializeCustomizablePermissionlessLbPairDiscriminator,
			accountCount:        14,
			wantPoolPos:         0,
			wantBaseMintPos:     2,
			wantQuoteMintPos:    3,
			wantUserPos:         8,
			wantBaseProgramPos:  9,
			wantQuoteProgramPos: 9,
		},
		{
			name:                "initialize_customizable_permissionless_lb_pair2",
			discriminator:       meteoraInitializeCustomizablePermissionlessLbPair2Discriminator,
			accountCount:        17,
			wantPoolPos:         0,
			wantBaseMintPos:     2,
			wantQuoteMintPos:    3,
			wantUserPos:         8,
			wantBaseProgramPos:  11,
			wantQuoteProgramPos: 12,
		},
		{
			name:                "initialize_permission_lb_pair",
			discriminator:       meteoraInitializePermissionLbPairDiscriminator,
			accountCount:        17,
			wantPoolPos:         1,
			wantBaseMintPos:     3,
			wantQuoteMintPos:    4,
			wantUserPos:         8,
			wantBaseProgramPos:  11,
			wantQuoteProgramPos: 12,
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			accountList := make([]solana.PublicKey, 32)
			for i := range accountList {
				accountList[i] = testPublicKey(byte(i + 1))
			}
			programIndex := 30
			accountList[programIndex] = meteoraDlmmProgram

			instruction := Instruction{
				Accounts:       seqInts(tc.accountCount),
				Data:           solana.Base58(tc.discriminator[:]),
				ProgramIDIndex: programIndex,
			}

			rawTx := &RawTx{
				accountList: accountList,
				Meta: Meta{
					PostTokenBalances: []TokenBalance{
						{
							MintAccount: accountList[tc.wantBaseMintPos],
							UITokenAmount: UITokenAmount{
								Decimals: 6,
							},
						},
						{
							MintAccount: accountList[tc.wantQuoteMintPos],
							UITokenAmount: UITokenAmount{
								Decimals: 9,
							},
						},
					},
				},
				Transaction: Transaction{
					Message: Message{
						Instructions: []Instruction{instruction},
					},
				},
			}

			tx := &Tx{rawTx: rawTx}

			swaps, _, err := metaoradlmmParser(tx, instruction, InnerInstructions{}, [2]uint{0, 0})
			if err != nil {
				t.Fatalf("metaoradlmmParser() error = %v", err)
			}
			if len(swaps) != 1 {
				t.Fatalf("metaoradlmmParser() swaps len = %d, want 1", len(swaps))
			}

			swap := swaps[0]
			if !swap.Pool.Equals(accountList[tc.wantPoolPos]) {
				t.Fatalf("swap.Pool = %s, want %s", swap.Pool, accountList[tc.wantPoolPos])
			}
			if !swap.BaseMint.Equals(accountList[tc.wantBaseMintPos]) {
				t.Fatalf("swap.BaseMint = %s, want %s", swap.BaseMint, accountList[tc.wantBaseMintPos])
			}
			if !swap.QuoteMint.Equals(accountList[tc.wantQuoteMintPos]) {
				t.Fatalf("swap.QuoteMint = %s, want %s", swap.QuoteMint, accountList[tc.wantQuoteMintPos])
			}
			if !swap.User.Equals(accountList[tc.wantUserPos]) {
				t.Fatalf("swap.User = %s, want %s", swap.User, accountList[tc.wantUserPos])
			}
			if !swap.BaseTokenProgram.Equals(accountList[tc.wantBaseProgramPos]) {
				t.Fatalf("swap.BaseTokenProgram = %s, want %s", swap.BaseTokenProgram, accountList[tc.wantBaseProgramPos])
			}
			if !swap.QuoteTokenProgram.Equals(accountList[tc.wantQuoteProgramPos]) {
				t.Fatalf("swap.QuoteTokenProgram = %s, want %s", swap.QuoteTokenProgram, accountList[tc.wantQuoteProgramPos])
			}
			if swap.BaseMintDecimals != 6 {
				t.Fatalf("swap.BaseMintDecimals = %d, want 6", swap.BaseMintDecimals)
			}
			if swap.QuoteMintDecimals != 9 {
				t.Fatalf("swap.QuoteMintDecimals = %d, want 9", swap.QuoteMintDecimals)
			}
			if !swap.EntryContract.Equals(meteoraDlmmProgram) {
				t.Fatalf("swap.EntryContract = %s, want %s", swap.EntryContract, meteoraDlmmProgram)
			}
		})
	}
}

func TestDlmmDecodeLbPairCreateEvent(t *testing.T) {
	t.Parallel()

	event := dlmmLbPairCreateEvent{
		LbPair:  testPublicKey(90),
		BinStep: 42,
		TokenX:  testPublicKey(91),
		TokenY:  testPublicKey(92),
	}

	body := mustBorshEncode(t, event)

	barePayload := append(append([]byte{}, meteoraInitializeLbPairEventDiscriminator[:]...), body...)
	decodedBare, ok := dlmmDecodeLbPairCreateEvent(barePayload)
	if !ok {
		t.Fatalf("dlmmDecodeLbPairCreateEvent() failed for bare payload")
	}
	if decodedBare != event {
		t.Fatalf("decoded bare event = %+v, want %+v", decodedBare, event)
	}

	anchorPayload := append(append(append([]byte{}, eventDiscriminator[:]...), meteoraInitializeLbPairEventDiscriminator[:]...), body...)
	decodedAnchor, ok := dlmmDecodeLbPairCreateEvent(anchorPayload)
	if !ok {
		t.Fatalf("dlmmDecodeLbPairCreateEvent() failed for anchor payload")
	}
	if decodedAnchor != event {
		t.Fatalf("decoded anchor event = %+v, want %+v", decodedAnchor, event)
	}
}

func TestResolveDlmmSwapAccountsAllowsRemainingAccountsAfterEventAuthority(t *testing.T) {
	t.Parallel()

	accountList := make([]solana.PublicKey, 40)
	for i := range accountList {
		accountList[i] = testPublicKey(byte(i + 1))
	}
	accountList[0] = testPublicKey(200)
	accountList[26] = meteoraDlmmProgram
	accountList[27] = solana.MemoProgramID
	accountList[29] = solana.TokenProgramID
	accountList[33] = meteoraDlmmEventAuthority

	rawTx := &RawTx{
		accountList: accountList,
		Transaction: Transaction{
			Message: Message{
				AccountKeys: accountList[:11],
				Header: Header{
					NumRequiredSignatures: 1,
				},
			},
		},
	}
	accounts := []int{13, 26, 16, 14, 11, 4, 35, 28, 15, 26, 0, 29, 29, 27, 33, 29, 3, 7, 2}

	resolved, err := resolveDlmmSwapAccounts(rawTx, accounts)
	if err != nil {
		t.Fatalf("resolveDlmmSwapAccounts() error = %v", err)
	}
	if resolved.poolIdx != 13 {
		t.Fatalf("poolIdx = %d, want 13", resolved.poolIdx)
	}
	if resolved.reserveXIdx != 16 || resolved.reserveYIdx != 14 {
		t.Fatalf("reserve indexes = %d/%d, want 16/14", resolved.reserveXIdx, resolved.reserveYIdx)
	}
	if resolved.userIdx != 0 {
		t.Fatalf("userIdx = %d, want 0", resolved.userIdx)
	}
	if resolved.tokenXProgramIdx != 29 || resolved.tokenYProgramIdx != 29 {
		t.Fatalf("token program indexes = %d/%d, want 29/29", resolved.tokenXProgramIdx, resolved.tokenYProgramIdx)
	}
}

func TestMeteoraDlmmInitializeParserUsesLbPairCreateEvent(t *testing.T) {
	t.Parallel()

	accountList := make([]solana.PublicKey, 32)
	for i := range accountList {
		accountList[i] = testPublicKey(byte(i + 1))
	}
	programIndex := 30
	accountList[programIndex] = meteoraDlmmProgram

	instruction := Instruction{
		Accounts:       seqInts(16),
		Data:           solana.Base58(meteoraInitializeLbPair2Discriminator[:]),
		ProgramIDIndex: programIndex,
	}

	event := dlmmLbPairCreateEvent{
		LbPair:  testPublicKey(111),
		BinStep: 25,
		TokenX:  testPublicKey(112),
		TokenY:  testPublicKey(113),
	}
	innerEventData := append(
		append(append([]byte{}, eventDiscriminator[:]...), meteoraInitializeLbPairEventDiscriminator[:]...),
		mustBorshEncode(t, event)...,
	)

	rawTx := &RawTx{
		accountList: accountList,
		Meta: Meta{
			PostTokenBalances: []TokenBalance{
				{
					MintAccount: accountList[2],
					UITokenAmount: UITokenAmount{
						Decimals: 6,
					},
				},
				{
					MintAccount: accountList[3],
					UITokenAmount: UITokenAmount{
						Decimals: 9,
					},
				},
			},
			InnerInstructions: []InnerInstructions{
				{
					Index: 0,
					Instructions: []Instruction{
						{
							ProgramIDIndex: programIndex,
							Data:           solana.Base58(innerEventData),
						},
					},
				},
			},
		},
		Transaction: Transaction{
			Message: Message{
				Instructions: []Instruction{instruction},
			},
		},
	}

	tx := &Tx{rawTx: rawTx}

	swaps, nextOffset, err := metaoradlmmParser(tx, instruction, rawTx.Meta.InnerInstructions[0], [2]uint{0, 0})
	if err != nil {
		t.Fatalf("metaoradlmmParser() error = %v", err)
	}
	if len(swaps) != 1 {
		t.Fatalf("metaoradlmmParser() swaps len = %d, want 1", len(swaps))
	}

	swap := swaps[0]
	if !swap.Pool.Equals(event.LbPair) {
		t.Fatalf("swap.Pool = %s, want event %s", swap.Pool, event.LbPair)
	}
	if !swap.BaseMint.Equals(event.TokenX) {
		t.Fatalf("swap.BaseMint = %s, want event %s", swap.BaseMint, event.TokenX)
	}
	if !swap.QuoteMint.Equals(event.TokenY) {
		t.Fatalf("swap.QuoteMint = %s, want event %s", swap.QuoteMint, event.TokenY)
	}
	if nextOffset != ([2]uint{1, 0}) {
		t.Fatalf("nextOffset = %#v, want [2]uint{1, 0}", nextOffset)
	}
}

func TestDlmmSwapFeeInfo(t *testing.T) {
	t.Parallel()

	baseMint := testPublicKey(1)
	quoteMint := testPublicKey(2)
	baseProgram := testPublicKey(3)
	quoteProgram := testPublicKey(4)

	testCases := []struct {
		name         string
		baseIsX      bool
		swapForY     bool
		wantFeeSide  string
		wantFeeMint  solana.PublicKey
		wantFeeProg  solana.PublicKey
		wantDecimals uint8
	}{
		{
			name:         "x is base and input is x",
			baseIsX:      true,
			swapForY:     true,
			wantFeeSide:  "base",
			wantFeeMint:  baseMint,
			wantFeeProg:  baseProgram,
			wantDecimals: 6,
		},
		{
			name:         "x is base and input is y",
			baseIsX:      true,
			swapForY:     false,
			wantFeeSide:  "quote",
			wantFeeMint:  quoteMint,
			wantFeeProg:  quoteProgram,
			wantDecimals: 9,
		},
		{
			name:         "y is base and input is x",
			baseIsX:      false,
			swapForY:     true,
			wantFeeSide:  "quote",
			wantFeeMint:  quoteMint,
			wantFeeProg:  quoteProgram,
			wantDecimals: 9,
		},
		{
			name:         "y is base and input is y",
			baseIsX:      false,
			swapForY:     false,
			wantFeeSide:  "base",
			wantFeeMint:  baseMint,
			wantFeeProg:  baseProgram,
			wantDecimals: 6,
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			feeAmount, feeSide, feeMint, feeProgram, feeDecimals := dlmmSwapFeeInfo(
				tc.baseIsX,
				tc.swapForY,
				123,
				baseMint,
				quoteMint,
				baseProgram,
				quoteProgram,
				6,
				9,
			)
			if !feeAmount.Equal(decimal.NewFromInt(123)) {
				t.Fatalf("feeAmount = %s, want 123", feeAmount)
			}
			if feeSide != tc.wantFeeSide {
				t.Fatalf("feeSide = %s, want %s", feeSide, tc.wantFeeSide)
			}
			if !feeMint.Equals(tc.wantFeeMint) {
				t.Fatalf("feeMint = %s, want %s", feeMint, tc.wantFeeMint)
			}
			if !feeProgram.Equals(tc.wantFeeProg) {
				t.Fatalf("feeProgram = %s, want %s", feeProgram, tc.wantFeeProg)
			}
			if feeDecimals != tc.wantDecimals {
				t.Fatalf("feeDecimals = %d, want %d", feeDecimals, tc.wantDecimals)
			}
		})
	}
}

func TestDlmmSwapLpFeeAmount(t *testing.T) {
	t.Parallel()

	lpFee := dlmmSwapLpFeeAmount(100, 15, 5)
	if !lpFee.Equal(decimal.NewFromInt(80)) {
		t.Fatalf("lpFee = %s, want 80", lpFee)
	}

	lpFee = dlmmSwapLpFeeAmount(10, 8, 5)
	if !lpFee.IsZero() {
		t.Fatalf("lpFee should floor at zero, got %s", lpFee)
	}
}

func TestDlmmSwapFeeBpsString(t *testing.T) {
	t.Parallel()

	feeBps := agbinary.Uint128{Lo: 12345}
	if got := dlmmSwapFeeBpsString(feeBps); got != "12345" {
		t.Fatalf("dlmmSwapFeeBpsString() = %s, want 12345", got)
	}
}