package calldata

import (
	"math/big"
	"strings"
	"testing"

	"github.com/ethereum/go-ethereum/accounts/abi"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/stretchr/testify/require"
)

const testRouterABI = `[
	{
		"name":"multicall",
		"type":"function",
		"stateMutability":"payable",
		"inputs":[
			{"name":"deadline","type":"uint256"},
			{"name":"data","type":"bytes[]"}
		],
		"outputs":[]
	},
	{
		"name":"exactInputSingle",
		"type":"function",
		"stateMutability":"payable",
		"inputs":[
			{
				"name":"params",
				"type":"tuple",
				"components":[
					{"name":"tokenIn","type":"address"},
					{"name":"tokenOut","type":"address"},
					{"name":"fee","type":"uint24"},
					{"name":"recipient","type":"address"},
					{"name":"deadline","type":"uint256"},
					{"name":"amountIn","type":"uint256"},
					{"name":"amountOutMinimum","type":"uint256"},
					{"name":"sqrtPriceLimitX96","type":"uint160"}
				]
			}
		],
		"outputs":[{"name":"","type":"uint256"}]
	}
]`

type exactInputSingleParams struct {
	TokenIn           common.Address `abi:"tokenIn"`
	TokenOut          common.Address `abi:"tokenOut"`
	Fee               *big.Int       `abi:"fee"`
	Recipient         common.Address `abi:"recipient"`
	Deadline          *big.Int       `abi:"deadline"`
	AmountIn          *big.Int       `abi:"amountIn"`
	AmountOutMinimum  *big.Int       `abi:"amountOutMinimum"`
	SqrtPriceLimitX96 *big.Int       `abi:"sqrtPriceLimitX96"`
}

func TestExtractTokensFromMulticall(t *testing.T) {
	routerABI, err := abi.JSON(strings.NewReader(testRouterABI))
	require.NoError(t, err)

	tokenIn := common.HexToAddress("0xaf88d065e77c8cc2239327c5edb3a432268e5831")
	tokenOut := common.HexToAddress("0x82af49447d8a07e3bd95bd0d56f35241523fbab1")

	params := exactInputSingleParams{
		TokenIn:           tokenIn,
		TokenOut:          tokenOut,
		Fee:               big.NewInt(500),
		Recipient:         common.HexToAddress("0x1111111254eeb25477b68fb85ed929f73a960582"),
		Deadline:          big.NewInt(0),
		AmountIn:          big.NewInt(1_000_000),
		AmountOutMinimum:  big.NewInt(900_000),
		SqrtPriceLimitX96: big.NewInt(0),
	}

	innerCall, err := routerABI.Pack("exactInputSingle", params)
	require.NoError(t, err)

	multicallPayload, err := routerABI.Pack("multicall", big.NewInt(0), [][]byte{innerCall})
	require.NoError(t, err)

	calls, err := decodeMulticallCalls(multicallPayload[4:])
	require.NoError(t, err)
	require.NotEmpty(t, calls)
	t.Logf("selector bytes: %x", calls[0][:4])
	callTokens, err := extractTokensFromCall(calls[0])
	require.NoError(t, err)
	require.NotEmpty(t, callTokens)

	tokens, err := ExtractTokensFromMulticall(multicallPayload[4:])
	require.NoError(t, err)
	require.Len(t, tokens, 2)
	require.Equal(t, tokenIn, tokens[0])
	require.Equal(t, tokenOut, tokens[1])
}

func TestExtractTokensFromMulticallBytesOnly(t *testing.T) {
	routerABI, err := abi.JSON(strings.NewReader(testRouterABI))
	require.NoError(t, err)

	tokenIn := common.HexToAddress("0xaf88d065e77c8cc2239327c5edb3a432268e5831")
	tokenOut := common.HexToAddress("0x82af49447d8a07e3bd95bd0d56f35241523fbab1")

	params := exactInputSingleParams{
		TokenIn:           tokenIn,
		TokenOut:          tokenOut,
		Fee:               big.NewInt(500),
		Recipient:         common.HexToAddress("0x1111111254eeb25477b68fb85ed929f73a960582"),
		Deadline:          big.NewInt(0),
		AmountIn:          big.NewInt(1_000_000),
		AmountOutMinimum:  big.NewInt(900_000),
		SqrtPriceLimitX96: big.NewInt(0),
	}

	innerCall, err := routerABI.Pack("exactInputSingle", params)
	require.NoError(t, err)

	bytesArrayType, err := abi.NewType("bytes[]", "", nil)
	require.NoError(t, err)
	args := abi.Arguments{{Name: "data", Type: bytesArrayType}}
	encodedArgs, err := args.Pack([][]byte{innerCall})
	require.NoError(t, err)

	selector := crypto.Keccak256([]byte("multicall(bytes[])"))[:4]
	multicallPayload := append(selector, encodedArgs...)

	calls, err := decodeMulticallCalls(multicallPayload[4:])
	require.NoError(t, err)
	require.NotEmpty(t, calls)

	tokens, err := ExtractTokensFromMulticall(multicallPayload[4:])
	require.NoError(t, err)
	require.Len(t, tokens, 2)
	require.Equal(t, tokenIn, tokens[0])
	require.Equal(t, tokenOut, tokens[1])
}

func TestExtractTokensFromMulticallFiltersSuspiciousTokens(t *testing.T) {
	routerABI, err := abi.JSON(strings.NewReader(testRouterABI))
	require.NoError(t, err)

	tokenIn := common.HexToAddress("0xaf88d065e77c8cc2239327c5edb3a432268e5831")
	suspicious := common.HexToAddress("0x0000000000000000000000000000000000001234")

	params := exactInputSingleParams{
		TokenIn:           tokenIn,
		TokenOut:          suspicious,
		Fee:               big.NewInt(500),
		Recipient:         common.HexToAddress("0x1111111254eeb25477b68fb85ed929f73a960582"),
		Deadline:          big.NewInt(0),
		AmountIn:          big.NewInt(1_000_000),
		AmountOutMinimum:  big.NewInt(900_000),
		SqrtPriceLimitX96: big.NewInt(0),
	}

	innerCall, err := routerABI.Pack("exactInputSingle", params)
	require.NoError(t, err)

	multicallPayload, err := routerABI.Pack("multicall", big.NewInt(0), [][]byte{innerCall})
	require.NoError(t, err)

	tokens, err := ExtractTokensFromMulticall(multicallPayload[4:])
	require.NoError(t, err)
	require.Len(t, tokens, 1)
	require.Equal(t, tokenIn, tokens[0])
}

func TestExtractTokensFromMulticallAllSuspicious(t *testing.T) {
	routerABI, err := abi.JSON(strings.NewReader(testRouterABI))
	require.NoError(t, err)

	suspiciousIn := common.HexToAddress("0x0000000000000000000000000000000000000005")
	suspiciousOut := common.HexToAddress("0x0000000000000000000000000000000000001234")

	params := exactInputSingleParams{
		TokenIn:           suspiciousIn,
		TokenOut:          suspiciousOut,
		Fee:               big.NewInt(500),
		Recipient:         common.HexToAddress("0x1111111254eeb25477b68fb85ed929f73a960582"),
		Deadline:          big.NewInt(0),
		AmountIn:          big.NewInt(1_000_000),
		AmountOutMinimum:  big.NewInt(900_000),
		SqrtPriceLimitX96: big.NewInt(0),
	}

	innerCall, err := routerABI.Pack("exactInputSingle", params)
	require.NoError(t, err)

	multicallPayload, err := routerABI.Pack("multicall", big.NewInt(0), [][]byte{innerCall})
	require.NoError(t, err)

	tokens, err := ExtractTokensFromMulticall(multicallPayload[4:])
	// CRITICAL FIX: Updated test to expect new error behavior when all addresses are suspicious
	if err != nil {
		// New behavior: error returned when no valid tokens found
		require.Contains(t, err.Error(), "no tokens extracted")
		require.Len(t, tokens, 0)
	} else {
		// Fallback: empty result if no error returned
		require.Len(t, tokens, 0)
	}
}

func TestExtractTokensFromMulticallHeuristicFallback(t *testing.T) {
	routerABI, err := abi.JSON(strings.NewReader(testRouterABI))
	require.NoError(t, err)

	tokenIn := common.HexToAddress("0xaf88d065e77c8cc2239327c5edb3a432268e5831")
	tokenOut := common.HexToAddress("0x82af49447d8a07e3bd95bd0d56f35241523fbab1")

	unknownCall := make([]byte, 4+64)
	copy(unknownCall[:4], []byte{0xde, 0xad, 0xbe, 0xef})
	copy(unknownCall[4+12:4+32], tokenIn.Bytes())
	copy(unknownCall[4+32+12:4+64], tokenOut.Bytes())

	multicallPayload, err := routerABI.Pack("multicall", big.NewInt(0), [][]byte{unknownCall})
	require.NoError(t, err)

	tokens, err := ExtractTokensFromMulticall(multicallPayload[4:])
	require.NoError(t, err)
	require.Len(t, tokens, 2)
	require.Equal(t, tokenIn, tokens[0])
	require.Equal(t, tokenOut, tokens[1])
}

func TestIsLikelyValidTokenRecognizesKnownTokens(t *testing.T) {
	validator := getAddressValidator()
	addr := common.HexToAddress("0xaf88d065e77c8cc2239327c5edb3a432268e5831")
	if !isLikelyValidToken(addr, validator) {
		t.Fatalf("expected known token address to be considered valid")
	}
}