import { type Aptos, AccountAddress } from '@aptos-labs/ts-sdk'
import { getBytes, zeroPadValue } from 'ethers'

import { encodeExtraArgs } from '../extra-args.ts'
import { type AnyMessage, ChainFamily } from '../types.ts'
import { getDataBytes } from '../utils.ts'

export const DEFAULT_FEE_TOKEN = '0xa'

function messageArgs(
  destChainSelector: bigint,
  message: AnyMessage,
): [
  destChainSelector: bigint,
  receiver: Uint8Array,
  data: Uint8Array,
  tokenAddresses: string[],
  tokenAmounts: string[],
  tokenStoreAddresses: string[],
  feeToken: string,
  feeTokenStore: string,
  encodedExtraArgs: Uint8Array,
] {
  // Prepare the message structure for the view call
  const receiver = getBytes(zeroPadValue(getDataBytes(message.receiver), 32))
  const data = getDataBytes(message.data || '0x')

  // Get the native token to use as fee token if not specified
  const feeToken = message.feeToken || DEFAULT_FEE_TOKEN
  const feeTokenStore = '0x0' // auto-fetch primary store

  // Split token amounts into separate arrays for Aptos Move
  const tokenAddresses = (message.tokenAmounts ?? []).map((ta) => ta.token)
  const tokenAmounts = (message.tokenAmounts ?? []).map((ta) => ta.amount.toString())
  const tokenStoreAddresses = tokenAddresses.map(() => '0x0')

  // Encode extraArgs for the router
  const encodedExtraArgs = getBytes(encodeExtraArgs(message.extraArgs, ChainFamily.Aptos))

  return [
    destChainSelector,
    receiver,
    data,
    tokenAddresses,
    tokenAmounts,
    tokenStoreAddresses,
    feeToken,
    feeTokenStore,
    encodedExtraArgs,
  ]
}

/**
 * Gets the fee for sending a CCIP message on Aptos.
 * @param provider - Aptos provider instance.
 * @param router - Router module address.
 * @param destChainSelector - Destination chain selector.
 * @param message - CCIP message to send.
 * @returns Fee amount in native tokens.
 */
export async function getFee(
  provider: Aptos,
  router: string,
  destChainSelector: bigint,
  message: AnyMessage,
): Promise<bigint> {
  // Call the get_fee view function on the router
  // Signature: get_fee(dest_chain_selector, receiver, data, token_addresses, token_amounts,
  //                    token_store_addresses, fee_token, fee_token_store, extra_args)
  const [fee] = await provider.view<[string]>({
    payload: {
      function:
        `${router.includes('::') ? router : router + '::router'}::get_fee` as `${string}::${string}::get_fee`,
      functionArguments: messageArgs(destChainSelector, message),
    },
  })

  return BigInt(fee)
}

/**
 * Generate an unsigned ccip_send tx to send a CCIP message through the Aptos router.
 * @param provider - Aptos provider instance.
 * @param sender - sender account address
 * @param router - Router module address.
 * @param destChainSelector - Destination chain selector.
 * @param message - CCIP message with fee.
 * @param _opts - Optional parameters.
 * @returns Transaction hash.
 */
export async function generateUnsignedCcipSend(
  provider: Aptos,
  sender: string,
  router: string,
  destChainSelector: bigint,
  message: AnyMessage & { fee: bigint },
  _opts?: { approveMax?: boolean },
): Promise<Uint8Array> {
  // Build and submit the transaction
  // Call ccip_send entry function with signature:
  // public entry fun ccip_send(
  //     caller: &signer,
  //     dest_chain_selector: u64,
  //     receiver: vector<u8>,
  //     data: vector<u8>,
  //     token_addresses: vector<address>,
  //     token_amounts: vector<u64>,
  //     token_store_addresses: vector<address>,
  //     fee_token: address,
  //     fee_token_store: address,
  //     extra_args: vector<u8>
  // )
  const transaction = await provider.transaction.build.simple({
    sender: AccountAddress.fromString(sender),
    data: {
      function:
        `${router.includes('::') ? router : router + '::router'}::ccip_send` as `${string}::${string}::${string}`,
      functionArguments: messageArgs(destChainSelector, message),
    },
  })

  return transaction.bcsToBytes()
}