const { ethers } = require("hardhat");
const fs = require("fs");

// 颜色输出
const colors = {
  reset: '\x1b[0m',
  red: '\x1b[31m',
  green: '\x1b[32m',
  yellow: '\x1b[33m',
  blue: '\x1b[34m',
  cyan: '\x1b[36m',
  white: '\x1b[37m'
};

function log(color, message) {
  console.log(`${colors[color]}${message}${colors.reset}`);
}

async function main() {
  log('cyan', '🚀 开始 MiniSwap AMM 完整测试场景');
  
  // 加载合约地址
  const addresses = JSON.parse(fs.readFileSync("contract-addresses.json", "utf8"));
  
  // 获取签名者
  const signers = await ethers.getSigners();
  const [deployer, user1, user2] = signers;
  
  // 连接到合约
  const TokenA = await ethers.getContractFactory("TokenA");
  const TokenB = await ethers.getContractFactory("TokenB");
  const MiniSwapAMM = await ethers.getContractFactory("MiniSwapAMM");
  
  const tokenA = TokenA.attach(addresses.TokenA);
  const tokenB = TokenB.attach(addresses.TokenB);
  const amm = MiniSwapAMM.attach(addresses.MiniSwapAMM);
  
  log('blue', '📄 合约连接成功');
  
  // 步骤 1: 显示初始状态
  log('yellow', '\n=== 步骤 1: 显示初始状态 ===');
  await showBalances('部署者', deployer, tokenA, tokenB, amm);
  await showBalances('用户1', user1, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 2: 为用户铸造代币
  log('yellow', '\n=== 步骤 2: 为用户铸造代币 ===');
  const mintAmount = ethers.parseEther("5000");
  
  await tokenA.connect(user1).faucet(mintAmount);
  await tokenB.connect(user1).faucet(mintAmount);
  await tokenA.connect(user2).faucet(mintAmount);
  await tokenB.connect(user2).faucet(mintAmount);
  
  log('green', '✅ 为用户1和用户2各铸造了 5000 TKA 和 5000 TKB');
  
  // 步骤 3: 授权代币
  log('yellow', '\n=== 步骤 3: 授权代币给 AMM 合约 ===');
  const approveAmount = ethers.parseEther("10000");
  const ammAddress = await amm.getAddress();
  
  await tokenA.connect(user1).approve(ammAddress, approveAmount);
  await tokenB.connect(user1).approve(ammAddress, approveAmount);
  await tokenA.connect(user2).approve(ammAddress, approveAmount);
  await tokenB.connect(user2).approve(ammAddress, approveAmount);
  
  log('green', '✅ 用户1和用户2已授权代币给 AMM 合约');
  
  // 步骤 4: 用户1添加初始流动性
  log('yellow', '\n=== 步骤 4: 用户1添加初始流动性 ===');
  const liquidityA = ethers.parseEther("1000");
  const liquidityB = ethers.parseEther("2000"); // 1:2 比例
  
  const tx1 = await amm.connect(user1).addLiquidity(
    liquidityA,
    liquidityB,
    liquidityA,
    liquidityB
  );
  await tx1.wait();
  
  log('green', '✅ 用户1添加了 1000 TKA + 2000 TKB 的流动性');
  
  await showBalances('用户1', user1, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 5: 用户2进行交换 A -> B
  log('yellow', '\n=== 步骤 5: 用户2进行交换 (A -> B) ===');
  const swapAmount = ethers.parseEther("100");
  
  // 计算预期输出
  const [reserveA, reserveB] = await amm.getReserves();
  const expectedOut = await amm.getAmountOut(swapAmount, reserveA, reserveB);
  
  log('blue', `输入: 100 TKA`);
  log('blue', `预期输出: ${ethers.formatEther(expectedOut)} TKB`);
  
  const tx2 = await amm.connect(user2).swapAForB(swapAmount, 0);
  await tx2.wait();
  
  log('green', '✅ 用户2完成交换');
  
  await showBalances('用户2', user2, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 6: 用户2进行反向交换 B -> A
  log('yellow', '\n=== 步骤 6: 用户2进行反向交换 (B -> A) ===');
  const swapBackAmount = ethers.parseEther("50");
  
  // 重新获取储备量
  const [newReserveA, newReserveB] = await amm.getReserves();
  const expectedOutBack = await amm.getAmountOut(swapBackAmount, newReserveB, newReserveA);
  
  log('blue', `输入: 50 TKB`);
  log('blue', `预期输出: ${ethers.formatEther(expectedOutBack)} TKA`);
  
  const tx3 = await amm.connect(user2).swapBForA(swapBackAmount, 0);
  await tx3.wait();
  
  log('green', '✅ 用户2完成反向交换');
  
  await showBalances('用户2', user2, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 7: 用户1添加更多流动性
  log('yellow', '\n=== 步骤 7: 用户1添加更多流动性 ===');
  
  // 根据当前比例添加流动性
  const [currentReserveA, currentReserveB] = await amm.getReserves();
  const addAmountA = ethers.parseEther("200");
  const addAmountB = (addAmountA * currentReserveB) / currentReserveA;
  
  log('blue', `添加 ${ethers.formatEther(addAmountA)} TKA + ${ethers.formatEther(addAmountB)} TKB`);
  
  const tx4 = await amm.connect(user1).addLiquidity(
    addAmountA,
    addAmountB,
    addAmountA * 99n / 100n, // 1% 滑点
    addAmountB * 99n / 100n
  );
  await tx4.wait();
  
  log('green', '✅ 用户1添加了更多流动性');
  
  await showBalances('用户1', user1, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 8: 用户1移除部分流动性
  log('yellow', '\n=== 步骤 8: 用户1移除部分流动性 ===');
  
  const lpBalance = await amm.balanceOf(user1.address);
  const removeAmount = lpBalance / 4n; // 移除 1/4 的流动性
  
  log('blue', `移除 ${ethers.formatEther(removeAmount)} LP 代币`);
  
  const tx5 = await amm.connect(user1).removeLiquidity(
    removeAmount,
    0, // 简化起见，最小值设为0
    0
  );
  await tx5.wait();
  
  log('green', '✅ 用户1移除了部分流动性');
  
  await showBalances('用户1', user1, tokenA, tokenB, amm);
  await showReserves(amm);
  
  // 步骤 9: 多次小额交换测试价格影响
  log('yellow', '\n=== 步骤 9: 多次小额交换测试价格影响 ===');
  
  for (let i = 1; i <= 3; i++) {
    const smallSwap = ethers.parseEther("10");
    const tx = await amm.connect(user2).swapAForB(smallSwap, 0);
    await tx.wait();
    
    const [rA, rB] = await amm.getReserves();
    const price = Number(ethers.formatEther(rB)) / Number(ethers.formatEther(rA));
    log('blue', `第${i}次交换后，价格 (TKB/TKA): ${price.toFixed(6)}`);
  }
  
  // 最终状态
  log('yellow', '\n=== 最终状态 ===');
  await showBalances('用户1', user1, tokenA, tokenB, amm);
  await showBalances('用户2', user2, tokenA, tokenB, amm);
  await showReserves(amm);
  
  log('cyan', '\n🎉 测试场景完成！');
  
  // 显示费用收益
  log('yellow', '\n=== 交易费用分析 ===');
  const finalReserves = await amm.getReserves();
  log('white', '由于每次交换都有 0.3% 的手续费');
  log('white', '这些费用会留在流动性池中，增加流动性提供者的收益');
  log('white', `最终储备量已经包含了累积的交易费用`);
}

async function showBalances(name, signer, tokenA, tokenB, amm) {
  const address = signer.address;
  const ethBalance = await signer.provider.getBalance(address);
  const tokenABalance = await tokenA.balanceOf(address);
  const tokenBBalance = await tokenB.balanceOf(address);
  const lpBalance = await amm.balanceOf(address);
  
  log('cyan', `\n--- ${name} 余额 (${address.slice(0,8)}...) ---`);
  log('white', `ETH: ${ethers.formatEther(ethBalance).slice(0,8)}`);
  log('white', `TKA: ${ethers.formatEther(tokenABalance)}`);
  log('white', `TKB: ${ethers.formatEther(tokenBBalance)}`);
  log('white', `LP:  ${ethers.formatEther(lpBalance)}`);
}

async function showReserves(amm) {
  const [reserveA, reserveB] = await amm.getReserves();
  const totalSupply = await amm.totalSupply();
  
  log('cyan', '\n--- AMM 储备量 ---');
  log('white', `TKA 储备: ${ethers.formatEther(reserveA)}`);
  log('white', `TKB 储备: ${ethers.formatEther(reserveB)}`);
  log('white', `LP 总量:  ${ethers.formatEther(totalSupply)}`);
  
  if (reserveA > 0 && reserveB > 0) {
    const price = Number(ethers.formatEther(reserveB)) / Number(ethers.formatEther(reserveA));
    log('white', `价格 (TKB/TKA): ${price.toFixed(6)}`);
  }
}

main()
  .then(() => process.exit(0))
  .catch((error) => {
    console.error(error);
    process.exit(1);
  });