approve.rs

#![allow(clippy::exhaustive_enums, reason = "Generated by sol! macro")]
#![allow(clippy::exhaustive_structs, reason = "Generated by sol! macro")]

use alloy::primitives::U256;
use alloy::sol;
use alloy::sol_types::SolCall;
use anyhow::{Context, Result};
use clap::{Args, Subcommand};
use polymarket_client_sdk::types::{Address, address};
use polymarket_client_sdk::{POLYGON, contract_config};

use crate::auth;
use crate::output::OutputFormat;
use crate::output::approve::{ApprovalStatus, print_approval_status, print_tx_result};

use super::proxy;

/// Polygon USDC (same address as `USDC_ADDRESS_STR`; `address!` requires a literal).
const USDC_ADDRESS: Address = address!("0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174");

sol! {
    #[sol(rpc)]
    interface IERC20 {
        function approve(address spender, uint256 value) external returns (bool);
        function allowance(address owner, address spender) external view returns (uint256);
    }

    #[sol(rpc)]
    interface IERC1155 {
        function setApprovalForAll(address operator, bool approved) external;
        function isApprovedForAll(address account, address operator) external view returns (bool);
    }
}

#[derive(Args)]
pub struct ApproveArgs {
    #[command(subcommand)]
    pub command: ApproveCommand,
}

#[derive(Subcommand)]
pub enum ApproveCommand {
    /// Check current contract approvals for a wallet
    Check {
        /// Wallet address to check (defaults to configured wallet)
        address: Option<Address>,
    },
    /// Approve all required contracts for trading (sends on-chain transactions)
    Set,
}

struct ApprovalTarget {
    name: &'static str,
    address: Address,
}

fn approval_targets() -> Result<Vec<ApprovalTarget>> {
    let config = contract_config(POLYGON, false).context("No contract config for Polygon")?;
    let neg_risk_config =
        contract_config(POLYGON, true).context("No neg-risk contract config for Polygon")?;

    let mut targets = vec![
        ApprovalTarget {
            name: "CTF Exchange",
            address: config.exchange,
        },
        ApprovalTarget {
            name: "Neg Risk Exchange",
            address: neg_risk_config.exchange,
        },
    ];

    if let Some(adapter) = neg_risk_config.neg_risk_adapter {
        targets.push(ApprovalTarget {
            name: "Neg Risk Adapter",
            address: adapter,
        });
    }

    Ok(targets)
}

pub async fn execute(
    args: ApproveArgs,
    output: OutputFormat,
    private_key: Option<&str>,
    signature_type: Option<&str>,
) -> Result<()> {
    match args.command {
        ApproveCommand::Check { address } => {
            check(address, private_key, signature_type, output).await
        }
        ApproveCommand::Set => set(private_key, signature_type, output).await,
    }
}

async fn check(
    address_arg: Option<Address>,
    private_key: Option<&str>,
    signature_type: Option<&str>,
    output: OutputFormat,
) -> Result<()> {
    let owner: Address = if let Some(addr) = address_arg {
        addr
    } else if proxy::is_proxy_mode(signature_type)? {
        proxy::derive_proxy_address(private_key)?
    } else {
        let signer = auth::resolve_signer(private_key)?;
        polymarket_client_sdk::auth::Signer::address(&signer)
    };

    let provider = auth::create_readonly_provider().await?;
    let config = contract_config(POLYGON, false).context("No contract config for Polygon")?;

    let usdc = IERC20::new(USDC_ADDRESS, provider.clone());
    let ctf = IERC1155::new(config.conditional_tokens, provider.clone());

    let targets = approval_targets()?;
    let mut statuses = Vec::new();

    for target in &targets {
        let (usdc_allowance, usdc_error) = match usdc.allowance(owner, target.address).call().await
        {
            Ok(val) => (val, None),
            Err(e) => (U256::ZERO, Some(e.to_string())),
        };

        let (ctf_approved, ctf_error) =
            match ctf.isApprovedForAll(owner, target.address).call().await {
                Ok(val) => (val, None),
                Err(e) => (false, Some(e.to_string())),
            };

        statuses.push(ApprovalStatus {
            contract_name: target.name.to_string(),
            contract_address: format!("{}", target.address),
            usdc_allowance,
            ctf_approved,
            usdc_error,
            ctf_error,
        });
    }

    print_approval_status(&statuses, &output)
}

async fn set(
    private_key: Option<&str>,
    signature_type: Option<&str>,
    output: OutputFormat,
) -> Result<()> {
    let use_proxy = proxy::is_proxy_mode(signature_type)?;
    let config = contract_config(POLYGON, false).context("No contract config for Polygon")?;
    let targets = approval_targets()?;
    let total = targets.len() * 2;

    if matches!(output, OutputFormat::Table) {
        println!("Approving contracts...\n");
    }

    let mut results: Vec<serde_json::Value> = Vec::new();
    let mut step = 0;

    if use_proxy {
        for target in &targets {
            step += 1;
            let label = format!("USDC \u{2192} {}", target.name);
            let calldata = IERC20::approveCall {
                spender: target.address,
                value: U256::MAX,
            }
            .abi_encode();
            let tx_hash = proxy::send_via_factory(private_key, USDC_ADDRESS, calldata)
                .await
                .context(format!("Failed USDC approval for {}", target.name))?;

            match output {
                OutputFormat::Table => print_tx_result(step, total, &label, tx_hash),
                OutputFormat::Json => results.push(serde_json::json!({
                    "step": step,
                    "type": "erc20",
                    "contract": target.name,
                    "tx_hash": format!("{tx_hash}"),
                })),
            }

            step += 1;
            let label = format!("CTF  \u{2192} {}", target.name);
            let calldata = IERC1155::setApprovalForAllCall {
                operator: target.address,
                approved: true,
            }
            .abi_encode();
            let tx_hash = proxy::send_via_factory(private_key, config.conditional_tokens, calldata)
                .await
                .context(format!("Failed CTF approval for {}", target.name))?;

            match output {
                OutputFormat::Table => print_tx_result(step, total, &label, tx_hash),
                OutputFormat::Json => results.push(serde_json::json!({
                    "step": step,
                    "type": "erc1155",
                    "contract": target.name,
                    "tx_hash": format!("{tx_hash}"),
                })),
            }
        }
    } else {
        let provider = auth::create_provider(private_key).await?;
        let usdc = IERC20::new(USDC_ADDRESS, provider.clone());
        let ctf = IERC1155::new(config.conditional_tokens, provider.clone());

        for target in &targets {
            step += 1;
            let label = format!("USDC \u{2192} {}", target.name);
            let tx_hash = usdc
                .approve(target.address, U256::MAX)
                .send()
                .await
                .context(format!("Failed to send USDC approval for {}", target.name))?
                .watch()
                .await
                .context(format!(
                    "Failed to confirm USDC approval for {}",
                    target.name
                ))?;

            match output {
                OutputFormat::Table => print_tx_result(step, total, &label, tx_hash),
                OutputFormat::Json => results.push(serde_json::json!({
                    "step": step,
                    "type": "erc20",
                    "contract": target.name,
                    "tx_hash": format!("{tx_hash}"),
                })),
            }

            step += 1;
            let label = format!("CTF  \u{2192} {}", target.name);
            let tx_hash = ctf
                .setApprovalForAll(target.address, true)
                .send()
                .await
                .context(format!("Failed to send CTF approval for {}", target.name))?
                .watch()
                .await
                .context(format!(
                    "Failed to confirm CTF approval for {}",
                    target.name
                ))?;

            match output {
                OutputFormat::Table => print_tx_result(step, total, &label, tx_hash),
                OutputFormat::Json => results.push(serde_json::json!({
                    "step": step,
                    "type": "erc1155",
                    "contract": target.name,
                    "tx_hash": format!("{tx_hash}"),
                })),
            }
        }
    }

    match output {
        OutputFormat::Table => {
            println!("\nAll contracts approved. You're ready to trade.");
        }
        OutputFormat::Json => {
            println!("{}", serde_json::to_string_pretty(&results)?);
        }
    }

    Ok(())
}