import 'dotenv/config';

import axios from "axios";
import { HttpsProxyAgent } from "https-proxy-agent";
import { AssetType, Chain, ClobClient, SignatureTypeV2 } from "@polymarket/clob-client-v2";
import { deriveProxyWallet, RelayClient } from "@polymarket/builder-relayer-client";
import { BuilderConfig } from "@polymarket/builder-signing-sdk";
import {
  createPublicClient,
  createWalletClient,
  erc20Abi,
  formatUnits,
  http,
} from "viem";
import { privateKeyToAccount } from "viem/accounts";
import { polygon } from "viem/chains";

const CHAIN_ID = 137;
const HOST = "https://clob.polymarket.com";
const PUSD_ADDRESS = "0xC011a7E12a19f7B1f670d46F03B03f3342E82DFB";
const PUSD_DECIMALS = 6;
const PROXY_FACTORY_ADDRESS = "0xaB45c5A4B0c941a2F231C04C3f49182e1A254052";

const NODE_HTTP_PROXY = process.env.NODE_HTTP_PROXY;
const proxyAgent = NODE_HTTP_PROXY ? new HttpsProxyAgent(NODE_HTTP_PROXY) : undefined;
if (NODE_HTTP_PROXY) {
  axios.defaults.proxy = false;
  axios.defaults.httpAgent = proxyAgent;
  axios.defaults.httpsAgent = proxyAgent;
}

const getRequiredEnv = (key) => {
  const value = process.env[key];
  if (!value) {
    throw new Error(`${key} is required`);
  }
  return value;
}

const getOptionalEnv = (key) => {
  const value = process.env[key];
  return value && value.trim() ? value.trim() : undefined;
}

const normalizePrivateKey = (privateKey) => {
  return privateKey.startsWith("0x") ? privateKey : `0x${privateKey}`;
}

const getArgValue = (name) => {
  const prefix = `${name}=`;
  const arg = process.argv.slice(2).find(item => item.startsWith(prefix));
  return arg ? arg.slice(prefix.length) : undefined;
}

const normalizeWalletMode = (wallet) => {
  const value = (wallet || "both").toLowerCase();
  if (!["deposit", "proxy", "both"].includes(value)) {
    throw new Error("wallet must be deposit, proxy, or both");
  }
  return value;
}

const makeClobClient = ({ signer, creds, funderAddress, signatureType }) => {
  return new ClobClient({
    host: HOST,
    chain: Chain.POLYGON,
    signer,
    creds,
    signatureType,
    funderAddress,
    throwOnError: true,
  });
}

const getChainBalance = async ({ publicClient, address }) => {
  const balance = await publicClient.readContract({
    address: PUSD_ADDRESS,
    abi: erc20Abi,
    functionName: "balanceOf",
    args: [address],
  });
  return formatUnits(balance, PUSD_DECIMALS);
}

const getClobBalanceAllowance = async ({ signer, creds, funderAddress, signatureType }) => {
  const client = makeClobClient({ signer, creds, funderAddress, signatureType });
  return client.getBalanceAllowance({ asset_type: AssetType.COLLATERAL });
}

const getDepositWalletAddress = async ({ account, signer, builderConfig }) => {
  const configuredAddress = getOptionalEnv("POLYMARKET_DEPOSIT_WALLET_ADDRESS");
  if (configuredAddress) {
    return configuredAddress;
  }

  const relayerUrl = getOptionalEnv("POLYMARKET_RELAYER_URL") || "https://relayer-v2.polymarket.com";
  const relayer = new RelayClient(relayerUrl, CHAIN_ID, signer, builderConfig);
  if (proxyAgent) {
    relayer.httpClient.instance.defaults.proxy = false;
    relayer.httpClient.instance.defaults.httpAgent = proxyAgent;
    relayer.httpClient.instance.defaults.httpsAgent = proxyAgent;
  }
  return relayer.deriveDepositWalletAddress();
}

const walletMode = normalizeWalletMode(getArgValue("--wallet") || process.env.BALANCE_WALLET);
const rpcUrl = getOptionalEnv("POLYGON_RPC_URL");
const transport = rpcUrl ? http(rpcUrl) : http();
const account = privateKeyToAccount(normalizePrivateKey(getRequiredEnv("POLYMARKET_PRIVATE_KEY")));
const signer = createWalletClient({ account, chain: polygon, transport });
const publicClient = createPublicClient({ chain: polygon, transport });
const creds = {
  key: getRequiredEnv("POLYMARKET_API_KEY"),
  secret: getRequiredEnv("POLYMARKET_API_SECRET"),
  passphrase: getRequiredEnv("POLYMARKET_API_PASSPHRASE"),
};
const builderConfig = new BuilderConfig({
  localBuilderCreds: {
    key: getRequiredEnv("POLYMARKET_BUILDER_API_KEY"),
    secret: getRequiredEnv("POLYMARKET_BUILDER_SECRET"),
    passphrase: getRequiredEnv("POLYMARKET_BUILDER_PASS_PHRASE"),
  },
});

const wallets = [];
if (walletMode === "proxy" || walletMode === "both") {
  wallets.push({
    type: "proxy",
    address: getOptionalEnv("POLYMARKET_PROXY_WALLET_ADDRESS")
      || deriveProxyWallet(account.address, PROXY_FACTORY_ADDRESS),
    signatureType: SignatureTypeV2.POLY_PROXY,
  });
}
if (walletMode === "deposit" || walletMode === "both") {
  wallets.push({
    type: "deposit",
    address: await getDepositWalletAddress({ account, signer, builderConfig }),
    signatureType: SignatureTypeV2.POLY_1271,
  });
}

const results = [];
for (const wallet of wallets) {
  const [chainPusdBalance, clobBalanceAllowance] = await Promise.all([
    getChainBalance({ publicClient, address: wallet.address }),
    getClobBalanceAllowance({
      signer,
      creds,
      funderAddress: wallet.address,
      signatureType: wallet.signatureType,
    }),
  ]);
  results.push({
    type: wallet.type,
    owner: account.address,
    address: wallet.address,
    signatureType: SignatureTypeV2[wallet.signatureType],
    chainPusdBalance,
    clobBalanceAllowance,
  });
}

console.log(JSON.stringify(results, null, 2));