package shreder

import (
	"bytes"
	"encoding/binary"
	"fmt"

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

var (
	maestroProgramId               = solana.MustPublicKeyFromBase58("MaestroAAe9ge5HTc64VbBQZ6fP77pwvrhM8i1XWSAx")
	maestroMultiSwap2Discriminator = [8]byte{132, 9, 212, 45, 39, 113, 215, 54}
)

type MaestroMultiSwap2Route struct {
	DexID       uint8
	RouteKeyIdx uint8
}

// MaestroMultiSwap2Args is the decoded payload of multi_swap2 instruction.
// Payload layout (without 8-byte discriminator):
// amount_in(u64), min_amount_out(u64), route0_weight(u16), route_len(u32),
// routes(route_len * {dex_id(u8), reserved(u8), route_key_idx(u8)}), route_family(u8), route_flags(u8)
type MaestroMultiSwap2Args struct {
	HasDiscriminator bool
	AmountIn         uint64
	MinAmountOut     uint64
	SlippageBps      uint16
	RouteLen         uint32
	Routes           []MaestroMultiSwap2Route
	RouteFamily      uint8
	RouteFlags       uint8
	Extra            []byte
}

// decodeMaestroMultiSwap2Args decodes instruction bytes with or without the 8-byte multi_swap2 discriminator.
func decodeMaestroMultiSwap2Args(data []byte) (*MaestroMultiSwap2Args, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("empty data")
	}

	payload := data
	out := &MaestroMultiSwap2Args{}

	if len(data) >= len(maestroMultiSwap2Discriminator) && bytes.Equal(data[:8], maestroMultiSwap2Discriminator[:]) {
		out.HasDiscriminator = true
		payload = data[8:]
	}

	const minPayloadLen = 24 // fixed fields + route_family + route_flags when route_len==0
	if len(payload) < minPayloadLen {
		return nil, fmt.Errorf("payload too short: got %d, need at least %d", len(payload), minPayloadLen)
	}

	out.AmountIn = binary.LittleEndian.Uint64(payload[0:8])
	out.MinAmountOut = binary.LittleEndian.Uint64(payload[8:16])
	out.SlippageBps = binary.LittleEndian.Uint16(payload[16:18])
	out.RouteLen = binary.LittleEndian.Uint32(payload[18:22])

	needed := uint64(minPayloadLen) + uint64(out.RouteLen)*3
	if needed > uint64(len(payload)) {
		return nil, fmt.Errorf("payload too short for routes: got %d, need %d (route_len=%d)", len(payload), needed, out.RouteLen)
	}

	offset := 22
	out.Routes = make([]MaestroMultiSwap2Route, 0, out.RouteLen)
	for i := uint32(0); i < out.RouteLen; i++ {
		route := MaestroMultiSwap2Route{
			DexID:       payload[offset],
			RouteKeyIdx: payload[offset+2],
		}
		out.Routes = append(out.Routes, route)
		offset += 3
	}

	out.RouteFamily = payload[offset]
	out.RouteFlags = payload[offset+1]
	offset += 2

	if len(payload) > offset {
		out.Extra = append([]byte(nil), payload[offset:]...)
	}

	return out, nil
}

// maestroMultiSwap2DexName maps observed dex ids from MultiSwap2 routes.
func maestroMultiSwap2DexName(dexID uint8) string {
	switch dexID {
	case 0:
		return "RaydiumV4"
	case 1:
		return "MeteoraDLMM"
	case 3:
		return "RaydiumLaunchLab"
	case 4:
		return "PumpAMM"
	case 5:
		return "RaydiumCPMM"
	case 6:
		return "MeteoraAmmV2"
	case 7:
		return "RaydiumCLMM"
	case 8:
		return "OrcaWhirPool"
	case 9:
		return "MeteoraPools"
	case 10:
		return "MeteoraDynamicBondingCurve"
	default:
		return fmt.Sprintf("Unknown(%d)", dexID)
	}
}

func parseMaestroInstruction(tx VersionedTransaction, idx int) (TxSignalBatch, error) {
	if idx >= len(tx.Instructions) {
		return nil, fmt.Errorf("instruction index out of bounds")
	}
	ix := tx.Instructions[idx]
	if len(ix.Data) < 8 {
		return nil, nil
	}
	return parseMaestroInstructionDataAndAccounts(tx, ix.Data, ix.Accounts)
}

func parsePumpAMMMaestroInstructionDataAndAccounts(tx VersionedTransaction, ixAccounts []uint8, args *MaestroMultiSwap2Args) (TxSignalBatch, error) {
	var (
		event        string
		token0Amount uint64
		token1Amount uint64

		// pool      solana.PublicKey
		baseMint  solana.PublicKey
		quoteMint solana.PublicKey
	)

	routeFlag := args.Routes[0].RouteKeyIdx
	if routeFlag == 101 || routeFlag == 97 {
		event = "buy"
		token0Amount = args.MinAmountOut
		token1Amount = args.AmountIn
	} else if routeFlag == 65 || routeFlag == 71 {
		event = "sell"
		token0Amount = args.AmountIn
		token1Amount = args.MinAmountOut
	} else {
		return nil, nil
	}

	if routeFlag == 101 || routeFlag == 71 {
		if len(ixAccounts) < 22 {
			return nil, nil
		}
		token2022, err := tx.GetAccount(int(ixAccounts[6]))
		if err != nil {
			return nil, err
		}
		if !token2022.Equals(solana.Token2022ProgramID) {
			return nil, nil
		}
		//pool, err = tx.GetAccount(int(ixAccounts[10]))
		if event == "buy" {
			quoteMint, err = tx.GetAccount(int(ixAccounts[9]))
			if err != nil {
				return nil, err
			}
			baseMint, err = tx.GetAccount(int(ixAccounts[21]))
			if err != nil {
				return nil, err
			}

		} else {
			baseMint, err = tx.GetAccount(int(ixAccounts[9]))
			if err != nil {
				return nil, err
			}

			quoteMint, err = tx.GetAccount(int(ixAccounts[21]))
			if err != nil {
				return nil, err
			}

		}
	} else if routeFlag == 97 || routeFlag == 65 {
		if len(ixAccounts) < 21 {
			return nil, nil
		}
		tokenPro, err := tx.GetAccount(int(ixAccounts[5]))
		if err != nil {
			return nil, err
		}
		if !tokenPro.Equals(solana.TokenProgramID) {
			return nil, nil
		}
		//pool, err = tx.GetAccount(int(ixAccounts[9]))
		if event == "buy" {
			quoteMint, err = tx.GetAccount(int(ixAccounts[8]))
			if err != nil {
				return nil, err
			}

			baseMint, err = tx.GetAccount(int(ixAccounts[20]))
			if err != nil {
				return nil, err
			}

		} else {
			baseMint, err = tx.GetAccount(int(ixAccounts[8]))
			if err != nil {
				return nil, err
			}

			quoteMint, err = tx.GetAccount(int(ixAccounts[20]))
			if err != nil {
				return nil, err
			}

		}
	}
	if !quoteMint.Equals(wrappedSOL) {
		return nil, nil
	}

	return TxSignalBatch{
		&TxSignal{
			TxHash:             tx.Signatures[0].String(),
			Maker:              tx.StaticAccountKeys[0].String(),
			Token0Address:      baseMint.String(),
			Token1Address:      wsolMint,
			Token0Amount:       formatTokenAmount(token0Amount),
			Token1Amount:       formatTokenAmount(token1Amount),
			Event:              event,
			Program:            "PumpAMM",
			IsProcessed:        false,
			IsToken2022:        false,
			IsMayhemMode:       false,
			ExactSOL:           event == "buy",
			Token0AmountUint64: token0Amount,
			Token1AmountUint64: token1Amount,
		},
	}, nil
}

func parseRaydiumLaunchLabMaestroInstructionDataAndAccounts(tx VersionedTransaction, ixAccounts []uint8, args *MaestroMultiSwap2Args) (TxSignalBatch, error) {
	var (
		event        string
		token0Amount uint64
		token1Amount uint64

		// pool      solana.PublicKey
		baseMint  solana.PublicKey
		quoteMint solana.PublicKey

		err error
	)

	routeFlag := args.Routes[0].RouteKeyIdx
	if routeFlag == 0x0b {
		event = "buy"
		token0Amount = args.MinAmountOut
		token1Amount = args.AmountIn
		if len(ixAccounts) < 18 {
			return nil, nil
		}
		quoteMint, err = tx.GetAccount(int(ixAccounts[7]))
		if err != nil {
			return nil, err
		}
		baseMint, err = tx.GetAccount(int(ixAccounts[17]))
		if err != nil {
			return nil, err
		}
	} else if routeFlag == 0x08 {
		event = "sell"
		token0Amount = args.AmountIn
		token1Amount = args.MinAmountOut
		if len(ixAccounts) < 18 {
			return nil, nil
		}
		quoteMint, err = tx.GetAccount(int(ixAccounts[17]))
		if err != nil {
			return nil, err
		}
		baseMint, err = tx.GetAccount(int(ixAccounts[7]))
		if err != nil {
			return nil, err
		}
	} else {
		return nil, nil
	}

	if !quoteMint.Equals(wrappedSOL) {
		return nil, nil
	}

	return TxSignalBatch{
		&TxSignal{
			TxHash:             tx.Signatures[0].String(),
			Maker:              tx.StaticAccountKeys[0].String(),
			Token0Address:      baseMint.String(),
			Token1Address:      wsolMint,
			Token0Amount:       formatTokenAmount(token0Amount),
			Token1Amount:       formatTokenAmount(token1Amount),
			Event:              event,
			Program:            "RaydiumLaunchLab",
			ExactSOL:           event == "buy",
			Token0AmountUint64: token0Amount,
			Token1AmountUint64: token1Amount,
		},
	}, nil
}

func parseMaestroInstructionDataAndAccounts(tx VersionedTransaction, ixData []byte, ixAccounts []uint8) (TxSignalBatch, error) {
	args, err := decodeMaestroMultiSwap2Args(ixData)
	if err != nil {
		return nil, nil
	}
	// only decode 1 route
	if len(args.Routes) != 1 {
		return nil, nil
	}

	if args.Routes[0].DexID == 4 {
		return parsePumpAMMMaestroInstructionDataAndAccounts(tx, ixAccounts, args)
	} else if args.Routes[0].DexID == 3 {
		return parseRaydiumLaunchLabMaestroInstructionDataAndAccounts(tx, ixAccounts, args)
	}

	return nil, nil
}