aa-swap.ts

import * as dotenv from 'dotenv';
import axios from 'axios';
import { ethers } from 'ethers';
import { Hex, hashMessage } from 'viem';

dotenv.config();

const PARTICLE_API_URL = 'https://rpc.particle.network';
const CHAIN_ID = 42161; // Arbitrum mainnet
const RPC_URL = 'https://arb1.arbitrum.io/rpc';

const rpcUrl = () => PARTICLE_API_URL;
const payloadId = () => Math.floor(Math.random() * Math.pow(10, 3));

// Define type interfaces needed for SmartAccount
interface IEthereumProvider {
  request(args: RequestArguments): Promise<any>;
  on?(eventName: string, listener: (...args: any[]) => void): any;
}

interface PasskeyProvider {
  isPasskey: boolean;
  getPasskeyOption?: () => Promise<any>;
}

interface RequestArguments {
  method: string;
  params?: any[];
}

interface AccountContract {
  name: string;
  version: string;
}

interface SmartAccountConfig {
  projectId: string;
  clientKey: string;
  appId: string;
  aaOptions: {
    accountContracts: {
      [key: string]: {
        version: string;
        chainIds?: number[];
      }[];
    };
  };
}

interface Account {
  smartAccountAddress: string;
  isDeployed: boolean;
}

interface AccountConfig {
  name: string;
  version: string;
  ownerAddress: string;
  options?: {
    passkeyOption?: any;
  };
}

interface Transaction {
  to: string;
  data: string;
  value?: string;
}

interface FeeQuotesResponse {
  verificationGasLimit?: string;
  preVerificationGas?: string;
  maxFeePerGas?: string;
  maxPriorityFeePerGas?: string;
  paymasterAndData?: string;
  callGasLimit?: string;
}

interface UserOp {
  sender: string;
  nonce: string;
  initCode: string;
  callData: string;
  callGasLimit: string;
  verificationGasLimit: string;
  preVerificationGas: string;
  maxFeePerGas: string;
  maxPriorityFeePerGas: string;
  paymasterAndData: string;
  signature: string;
}

interface UserOpBundle {
  userOp: UserOp;
  userOpHash: string;
}

interface UserOpParams {
  tx: Transaction | Transaction[];
  feeQuote?: FeeQuotesResponse;
  tokenPaymasterAddress?: string;
}

type SendTransactionParams = UserOpBundle | UserOpParams;

interface SessionKey {
  // Simplified type - can be extended if needed
  [key: string]: any;
}

interface CreateSessionKeyOptions {
  // Simplified type - can be extended if needed
  [key: string]: any;
}

interface SessionKeySignerParams {
  // Simplified type - can be extended if needed
  [key: string]: any;
}

// Implement SmartAccount class locally
class SmartAccount {
  private connection;
  private smartAccountContract: AccountContract;

  constructor(
    public provider: IEthereumProvider & Partial<PasskeyProvider>,
    private config: SmartAccountConfig,
  ) {
    if (
      !this.config.projectId ||
      !this.config.clientKey ||
      !this.config.appId
    ) {
      throw new Error('invalid project config');
    }
    if (!this.config.aaOptions.accountContracts) {
      throw new Error('invalid AA contract config');
    }
    const name = Object.keys(this.config.aaOptions.accountContracts)[0];
    const version = this.config.aaOptions.accountContracts[name]?.[0]?.version;
    if (!name || !version) {
      throw new Error('invalid AA name or version');
    }
    this.smartAccountContract = {
      name,
      version,
    };

    this.connection = axios.create({
      baseURL: `${rpcUrl()}/evm-chain`,
      timeout: 60_000,
    });

    this.connection.interceptors.request.use((config) => {
      if (config?.data?.method) {
        config.baseURL = `${config.baseURL}${config.baseURL?.includes('?') ? '&' : '?'}method=${
          config?.data?.method
        }`;
      }
      return config;
    });
  }

  setSmartAccountContract(contract: AccountContract) {
    const accountContract =
      this.config.aaOptions.accountContracts[contract.name];
    if (
      !accountContract ||
      accountContract.length === 0 ||
      accountContract.every((item) => item.version !== contract.version)
    ) {
      throw new Error('Please configure the smart account contract first');
    }
    this.smartAccountContract = contract;
  }

  getChainId = async (): Promise<string> => {
    return await this.provider.request({ method: 'eth_chainId' });
  };

  getOwner = async (): Promise<string> => {
    const eoas = await this.provider.request({ method: 'eth_accounts' });
    return eoas[0];
  };

  signUserOpHash = async (userOpHash: Hex): Promise<string> => {
    let message = userOpHash;
    if (
      this.provider.isPasskey &&
      this.smartAccountContract.name !== 'COINBASE'
    ) {
      message = hashMessage({
        raw: userOpHash,
      });
    }
    const eoa = await this.getOwner();
    const signature = await this.provider.request({
      method: 'personal_sign',
      params: [message, eoa],
    });
    return signature;
  };

  private async getAccountConfig(): Promise<AccountConfig> {
    const accountContract =
      this.config.aaOptions.accountContracts[this.smartAccountContract.name];
    if (
      !accountContract ||
      accountContract.every(
        (item) => item.version !== this.smartAccountContract.version,
      )
    ) {
      throw new Error('Please configure the smart account contract first');
    }

    const ownerAddress = await this.getOwner();
    let passkeyOption;
    if (this.provider.isPasskey) {
      passkeyOption = await this.provider.getPasskeyOption?.();
    }

    return {
      name: this.smartAccountContract.name,
      version: this.smartAccountContract.version,
      ownerAddress,
      options: passkeyOption
        ? {
            passkeyOption,
          }
        : undefined,
    };
  }

  async getFeeQuotes(
    tx: Transaction | Transaction[],
  ): Promise<FeeQuotesResponse> {
    const accountConfig = await this.getAccountConfig();
    return this.sendRpc<FeeQuotesResponse>({
      method: 'particle_aa_getFeeQuotes',
      params: [accountConfig, Array.isArray(tx) ? tx : [tx]],
    });
  }

  async buildUserOperation({
    tx,
    feeQuote,
    tokenPaymasterAddress,
  }: UserOpParams): Promise<UserOpBundle> {
    const accountConfig = await this.getAccountConfig();
    return await this.sendRpc<UserOpBundle>({
      method: 'particle_aa_createUserOp',
      params: [
        accountConfig,
        Array.isArray(tx) ? tx : [tx],
        feeQuote,
        tokenPaymasterAddress,
      ].filter((val) => !!val),
    });
  }

  async signUserOperation({
    userOpHash,
    userOp,
  }: UserOpBundle): Promise<UserOp> {
    const signature = await this.signUserOpHash(userOpHash as Hex);
    return { ...userOp, signature };
  }

  async sendUserOperation({
    userOpHash,
    userOp,
  }: UserOpBundle): Promise<string> {
    const signedUserOp = await this.signUserOperation({ userOpHash, userOp });
    return this.sendSignedUserOperation(signedUserOp);
  }

  async sendSignedUserOperation(
    userOp: UserOp,
    signerParams?: SessionKeySignerParams,
  ): Promise<string> {
    const accountConfig = await this.getAccountConfig();
    return this.sendRpc<string>({
      method: 'particle_aa_sendUserOp',
      params: [accountConfig, userOp, signerParams],
    });
  }

  async sendTransaction(params: SendTransactionParams): Promise<string> {
    if (
      Object.prototype.hasOwnProperty.call(params, 'userOpHash') &&
      Object.prototype.hasOwnProperty.call(params, 'userOp')
    ) {
      const { userOpHash, userOp } = params as UserOpBundle;
      if (userOpHash && userOp) {
        return this.sendUserOperation({ userOpHash, userOp });
      }
    }

    const { tx, feeQuote, tokenPaymasterAddress } = params as UserOpParams;
    const userOpBundle = await this.buildUserOperation({
      tx,
      feeQuote,
      tokenPaymasterAddress,
    });
    return this.sendUserOperation(userOpBundle);
  }

  async getAccount(): Promise<Account> {
    const accountConfig = await this.getAccountConfig();
    const accounts = await this.sendRpc<Account[]>({
      method: 'particle_aa_getSmartAccount',
      params: [accountConfig],
    });
    return accounts[0];
  }

  async getAddress(): Promise<string> {
    let suffix = await this.getOwner();
    if (!suffix) {
      return '';
    }

    if (
      this.provider.isPasskey &&
      suffix === '0x0000000000000000000000000000000000000000'
    ) {
      // passkey
      const credentialId = (await this.provider.getPasskeyOption?.())
        ?.credentialId;
      if (credentialId) {
        suffix = credentialId;
      }
    }

    const accountConfig = await this.getAccountConfig();
    const localKey = `particle_${accountConfig.name}_${accountConfig.version}_${suffix}`;
    if (typeof window !== 'undefined' && localStorage) {
      const localAA = localStorage.getItem(localKey);
      if (localAA) {
        return localAA;
      }
    }

    const configKey = JSON.stringify(accountConfig);
    let accountPromise = loadAccountPromise.get(configKey);
    if (!accountPromise) {
      accountPromise = this.getAccount();
      loadAccountPromise.set(configKey, accountPromise);
    }

    try {
      const account = await accountPromise;
      const address = account.smartAccountAddress;
      if (typeof window !== 'undefined' && localStorage) {
        localStorage.setItem(localKey, address);
      }
      return address;
    } catch (error) {
      loadAccountPromise.delete(configKey);
      throw error;
    }
  }

  async isDeployed(): Promise<boolean> {
    const account = await this.getAccount();
    return account.isDeployed;
  }

  async deployWalletContract(): Promise<string> {
    return this.sendTransaction({
      tx: {
        to: '0x0000000000000000000000000000000000000000',
        data: '0x',
      },
    });
  }

  async sendRpc<T>(arg: RequestArguments): Promise<T> {
    const chainId = Number(await this.getChainId());
    const accountContract =
      this.config.aaOptions.accountContracts[this.smartAccountContract.name];
    const contractConfig = accountContract.find(
      (contract) => contract.version === this.smartAccountContract.version,
    );
    if (contractConfig?.chainIds?.length) {
      if (!contractConfig.chainIds.includes(chainId)) {
        throw new Error(`Invalid Chain: ${chainId}`);
      }
    }
    const response = await this.connection
      .post(
        '',
        {
          ...arg,
          id: payloadId(),
          jsonrpc: '2.0',
        },
        {
          params: {
            chainId,
            projectUuid: this.config.projectId,
            projectKey: this.config.clientKey,
          },
        },
      )
      .then((res) => res.data);
    if (response.error) {
      return Promise.reject(response.error);
    } else {
      return response.result;
    }
  }

  async createSessions(
    options: CreateSessionKeyOptions[],
  ): Promise<FeeQuotesResponse> {
    const accountConfig = await this.getAccountConfig();
    return await this.sendRpc<FeeQuotesResponse>({
      method: 'particle_aa_createSessions',
      params: [accountConfig, options],
    });
  }

  async validateSession(
    targetSession: SessionKey,
    sessions: SessionKey[],
  ): Promise<boolean> {
    const accountConfig = await this.getAccountConfig();
    return await this.sendRpc<boolean>({
      method: 'particle_aa_validateSession',
      params: [
        accountConfig,
        {
          sessions,
          targetSession: targetSession,
        },
      ],
    });
  }
}

const loadAccountPromise = new Map<string, Promise<Account>>();

// Initialize wallet, smart account and get user addresses
async function initializeWalletAndSmartAccount() {
  const privateKey = process.env.PRIVATE_KEY;
  const projectId = process.env.PARTICLE_PROJECT_ID;
  const clientKey = process.env.PARTICLE_SERVER_KEY;
  const appId = process.env.PARTICLE_APP_ID;

  if (!privateKey) {
    throw new Error(
      'PRIVATE_KEY not found in environment variables. Add it to your .env file',
    );
  }

  if (!projectId || !clientKey || !appId) {
    throw new Error(
      'PARTICLE_PROJECT_ID, PARTICLE_SERVER_KEY, or PARTICLE_APP_ID not found in environment variables. Add them to your .env file',
    );
  }

  const provider = new ethers.JsonRpcProvider(RPC_URL);
  const wallet = new ethers.Wallet(privateKey, provider);

  const eip1193Provider = {
    request: async ({ method, params }: { method: string; params: any[] }) => {
      if (method === 'eth_accounts' || method === 'eth_requestAccounts') {
        return [wallet.address];
      }

      if (method === 'personal_sign') {
        const messageHex = params[0];
        const address = params[1];

        // Verify the address matches
        if (address.toLowerCase() !== wallet.address.toLowerCase()) {
          throw new Error('Address mismatch');
        }

        // Sign the message with the wallet's private key
        const signature = await wallet.signMessage(ethers.getBytes(messageHex));
        return signature;
      }

      return provider.send(method, params || []);
    },
    on: (eventName: string, listener: any) => {
      if (eventName === 'accountsChanged') {
        setTimeout(() => listener([wallet.address]), 0);
      }
      return eip1193Provider;
    },
  };

  const smartAccount = new SmartAccount(eip1193Provider, {
    projectId,
    clientKey,
    appId,
    aaOptions: {
      accountContracts: {
        BICONOMY: [
          {
            version: '2.0.0',
            chainIds: [CHAIN_ID],
          },
        ],
      },
    },
  });

  smartAccount.setSmartAccountContract({ name: 'BICONOMY', version: '2.0.0' });

  const eoaAddress = wallet.address;

  const smartAccountAddress = await smartAccount.getAddress();

  console.log('EOA Address:', eoaAddress);
  console.log('Smart Account Address:', smartAccountAddress);

  return { provider, wallet, smartAccount, eoaAddress, smartAccountAddress };
}

// Parse command line arguments
function parseArgs() {
  const args = process.argv.slice(2);
  const swapExactAmountOut = args.includes('--swapexactamountout');
  const swapExactAmountIn =
    args.includes('--swapexactamountin') || !swapExactAmountOut;

  return {
    swapExactAmountIn,
    swapExactAmountOut,
  };
}

function getParticleHeaders() {
  const projectId = process.env.PARTICLE_PROJECT_ID;
  const serverKey = process.env.PARTICLE_SERVER_KEY;

  if (!projectId || !serverKey) {
    throw new Error(
      'PARTICLE_PROJECT_ID or PARTICLE_SERVER_KEY not found in environment variables. Add them to your .env file',
    );
  }

  const auth = Buffer.from(`${projectId}:${serverKey}`).toString('base64');

  return {
    Authorization: `Basic ${auth}`,
    'Content-Type': 'application/json',
  };
}

async function particleRpcCall(
  method: string,
  params: any[],
  chainId = CHAIN_ID,
) {
  try {
    const response = await axios.post(
      `${PARTICLE_API_URL}/evm-chain`,
      {
        jsonrpc: '2.0',
        method,
        params,
        id: 1, // Added ID parameter for JSON-RPC
        chainId,
      },
      { headers: getParticleHeaders() },
    );

    if (response.data.error) {
      throw new Error(
        `Particle API Error: ${JSON.stringify(response.data.error)}`,
      );
    }

    return response.data.result;
  } catch (error) {
    if (axios.isAxiosError(error)) {
      console.error('API Response:', error.response?.data);
      throw new Error(`Network Error: ${error.message}`);
    }
    throw error;
  }
}

async function checkApprove(
  userAddress: string,
  token: { tokenAddress: string; amount: string },
  chainId = CHAIN_ID,
) {
  console.log('Checking token approval status...');

  // Native tokens don't need approval
  if (
    token.tokenAddress.toLowerCase() ===
    '0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee'
  ) {
    console.log('Token is native, no approval needed');
    return { approved: true };
  }

  const result = await particleRpcCall(
    'particle_swap_checkApprove',
    [userAddress, token],
    chainId,
  );

  console.log(
    `Approval status: ${result.approved ? 'Approved' : 'Not approved'}`,
  );
  return result;
}

async function executeApprovalTransaction(tx: any, smartAccount: SmartAccount) {
  console.log('Executing approval transaction using Smart Account...');

  try {
    const txParams = {
      tx: {
        to: tx.to,
        data: tx.data,
        value: tx.value || '0x0',
      },
    };

    console.log('Building UserOp for approval transaction...');
    const userOpBundle = await smartAccount.buildUserOperation(txParams);

    console.log('Sending UserOp for approval transaction...');
    const txHash = await smartAccount.sendUserOperation({
      userOp: userOpBundle.userOp,
      userOpHash: userOpBundle.userOpHash,
    });

    console.log(`Approval transaction sent! Hash: ${txHash}`);

    return { hash: txHash };
  } catch (error) {
    console.error('Failed to send approval transaction:', error);
    throw new Error(
      `Approval transaction failed: ${error instanceof Error ? error.message : String(error)}`,
    );
  }
}

async function getQuote(
  userAddress: string,
  fromToken: { tokenAddress: string; amount: string },
  toToken: { tokenAddress: string; amount?: string },
  swapMethod: string = 'exactAmountIn',
  chainId = CHAIN_ID,
) {
  console.log(`Getting swap quote for ${swapMethod}...`);

  try {
    const params: any = {
      fromTokenAddress: fromToken.tokenAddress,
      toTokenAddress: toToken.tokenAddress,
      slippage: 1, // 1% slippage as a number, not a string
    };

    // Set amount based on swap method
    if (swapMethod === 'exactAmountIn') {
      params.amount = fromToken.amount;
    } else if (swapMethod === 'exactAmountOut') {
      params.amount = toToken.amount;
      params.exactOut = true;
    }

    const result = await particleRpcCall(
      'particle_swap_getQuote',
      [userAddress, params],
      chainId,
    );

    console.log(
      `Quote received: ${fromToken.amount} -> ${result.toTokenAmount}`,
    );
    return result;
  } catch (error) {
    console.error(`Error getting quote: ${error.message}`);
    throw error;
  }
}

async function getSwap(
  userAddress: string,
  fromToken: { tokenAddress: string; amount: string },
  toToken: { tokenAddress: string; amount?: string },
  swapMethod: string = 'exactAmountIn',
  chainId = CHAIN_ID,
) {
  console.log(`Generating swap transaction for ${swapMethod}...`);

  const params: any = {
    fromTokenAddress: fromToken.tokenAddress,
    toTokenAddress: toToken.tokenAddress,
    slippage: 1, // 1% slippage as a number, not a string
  };

  // Set amount based on swap method
  if (swapMethod === 'exactAmountIn') {
    params.amount = fromToken.amount;
  } else if (swapMethod === 'exactAmountOut') {
    params.amount = toToken.amount;
    params.exactOut = true;
  }

  const result = await particleRpcCall(
    'particle_swap_getSwap',
    [userAddress, params],
    chainId,
  );

  console.log('Swap transaction generated');
  return result;
}

async function executeSwapTransaction(tx: any, smartAccount: SmartAccount) {
  console.log('Executing swap transaction using Smart Account...');

  try {
    const txParams = {
      tx: {
        to: tx.to,
        data: tx.data,
        value: tx.value || '0x0',
      },
    };

    console.log('Building UserOp for swap transaction...');
    const userOpBundle = await smartAccount.buildUserOperation(txParams);

    console.log('Sending UserOp for swap transaction...');
    const txHash = await smartAccount.sendUserOperation({
      userOp: userOpBundle.userOp,
      userOpHash: userOpBundle.userOpHash,
    });

    console.log(`Swap transaction sent! Hash: ${txHash}`);

    return {
      status: 'TRANSACTION_COMPLETE',
      hash: txHash,
    };
  } catch (error) {
    console.error('Failed to send swap transaction:', error);
    throw new Error(
      `Swap transaction failed: ${error instanceof Error ? error.message : String(error)}`,
    );
  }
}

async function ensureWalletDeployed(smartAccount: SmartAccount) {
  console.log('Checking if smart account is deployed...');
  const isDeployed = await smartAccount.isDeployed();

  if (!isDeployed) {
    console.log('Smart account not deployed. Deploying now...');
    const txHash = await smartAccount.deployWalletContract();
    console.log(`Smart account deployed! Hash: ${txHash}`);
    return txHash;
  }

  console.log('Smart account already deployed');
  return null;
}

async function executeSwap(
  config: {
    userAddress: string;
    fromToken: { tokenAddress: string; amount: string };
    toToken: { tokenAddress: string; amount?: string };
    swapMethod: string;
  },
  smartAccount: SmartAccount,
  chainId = CHAIN_ID,
) {
  try {
    const { userAddress, fromToken, toToken, swapMethod } = config;

    if (!userAddress) {
      throw new Error('User wallet address not available');
    }

    console.log('Starting swap process...');
    console.log(`Chain ID: ${chainId}`);
    console.log(`Swap Method: ${swapMethod}`);
    console.log(`From: ${fromToken.tokenAddress}`);
    console.log(`To: ${toToken.tokenAddress}`);

    if (swapMethod === 'exactAmountIn') {
      console.log(`Input Amount: ${fromToken.amount}`);
    } else if (swapMethod === 'exactAmountOut') {
      console.log(`Output Amount: ${toToken.amount}`);
    }

    console.log(`User address: ${userAddress}`);

    // Step 1: Check if token approval is needed
    const approvalStatus = await checkApprove(userAddress, fromToken, chainId);

    // Step 2: If approval is needed, execute approval transaction first
    if (!approvalStatus.approved && approvalStatus.tx) {
      console.log(
        '\nToken approval required. Executing approval transaction...',
      );
      await executeApprovalTransaction(approvalStatus.tx, smartAccount);
      console.log('Approval complete. Continuing with swap...');
    }

    // Step 3: Get swap quote
    const swapQuote = await getQuote(
      userAddress,
      fromToken,
      toToken,
      swapMethod,
      chainId,
    );

    if (swapMethod === 'exactAmountIn') {
      console.log(`\nExpected output amount: ${swapQuote.toTokenAmount}`);
    } else if (swapMethod === 'exactAmountOut') {
      console.log(`\nRequired input amount: ${swapQuote.fromTokenAmount}`);
      fromToken.amount = swapQuote.fromTokenAmount;
    }

    // Step 4: Get swap transaction
    const swapData = await getSwap(
      userAddress,
      fromToken,
      toToken,
      swapMethod,
      chainId,
    );

    // Step 5: Execute the swap
    const result = await executeSwapTransaction(swapData.tx, smartAccount);

    console.log('\nSwap completed successfully!');
    return result;
  } catch (error) {
    console.error('Error executing swap:', error);
    throw error;
  }
}

// Main function to handle the entire process
async function main() {
  // Initialize wallet and smart account
  const { wallet, smartAccount } = await initializeWalletAndSmartAccount();
  const userAddress = wallet.address;

  // Get command line arguments
  const args = parseArgs();
  console.log(
    `Running in ${args.swapExactAmountOut ? 'exact amount out' : 'exact amount in'} mode`,
  );

  // Create swap configuration based on command line args
  const swapConfig = {
    userAddress,
    fromToken: {
      tokenAddress: '0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', // Native token (ETH)
      amount: '100000000000000', // 0.0001 ETH in wei
    },
    toToken: {
      tokenAddress: '0xaf88d065e77c8cc2239327c5edb3a432268e5831', // USDC
      amount: '1000000', // 1 USDC (6 decimals)
    },
    swapMethod: args.swapExactAmountOut ? 'exactAmountOut' : 'exactAmountIn',
  };

  // Execute the swap with the configuration
  return executeSwap(swapConfig, smartAccount);
}

// Entry point when run directly
if (require.main === module) {
  main()
    .then(() => process.exit(0))
    .catch((error) => {
      console.error('Fatal error:', error);
      process.exit(1);
    });
}

// Create default swap config for exports
const DEFAULT_SWAP_CONFIG = {
  userAddress: '', // Will be filled in by whoever imports
  fromToken: {
    tokenAddress: '0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee', // Native token (ETH)
    amount: '100000000000000', // 0.0001 ETH in wei
  },
  toToken: {
    tokenAddress: '0xaf88d065e77c8cc2239327c5edb3a432268e5831', // USDC
    amount: '1000000', // 1 USDC (6 decimals)
  },
  swapMethod: 'exactAmountIn',
};

export {
  checkApprove,
  getQuote,
  getSwap,
  executeSwapTransaction,
  executeApprovalTransaction,
  executeSwap,
  initializeWalletAndSmartAccount,
  main,
  DEFAULT_SWAP_CONFIG,
};