main.py

"""
main.py — 端到端演示脚本

用法：
  # Balancer V2 验证（默认）
  python main.py --rpc https://eth-mainnet.g.alchemy.com/v2/YOUR_KEY
  python main.py --rpc YOUR_RPC_URL --pool 80BAL-20WETH
  python main.py --rpc YOUR_RPC_URL --scan

  # Uniswap V3 验证
  python main.py --rpc YOUR_RPC_URL --protocol v3 --pool USDC-WETH-3000

  # 列出所有已注册协议
  python main.py --list-protocols

环境变量方式（推荐使用 .env 存放密钥）：
  在 .env 中设置 ETH_RPC_URL=https://...
  python main.py
"""
import sys
import os
import argparse

from dotenv import load_dotenv

load_dotenv()

from web3 import Web3

from src import (
    ProtocolRegistry,
    BalancerV2Adapter,
    UniswapV3Adapter,
    UniswapV4Adapter,
    AlgebraAdapter,
    # Balancer 专用
    BalancerPoolMeta,
    KNOWN_POOLS,
    KNOWN_TOKENS,
    ONE,
    pool_meta_registry,
    pool_id_to_address,
    fetch_raw_state,
    validate_and_convert,
    validate_constraint_accuracy,
    check_pool_support_status,
    scan_all_pools,
    normalize_pool_id,
    # Event Scan
    scan_events_last_n_days,
    build_weekly_report,
    write_scan_outputs,
    # V3 专用
    UniswapV3PoolMeta,
)
from src.protocols.uniswap_v3.models import KNOWN_V3_POOLS, FEE_TIER_TICK_SPACING


# ── 协议注册表 ────────────────────────────────────────────

def build_registry() -> ProtocolRegistry:
    registry = ProtocolRegistry()
    registry.register(BalancerV2Adapter())
    registry.register(UniswapV3Adapter())
    registry.register(UniswapV4Adapter())
    algebra_chain = os.environ.get("ALGEBRA_CHAIN", "ethereum")
    registry.register(AlgebraAdapter(chain=algebra_chain))
    return registry


# ── Balancer V2 演示 ──────────────────────────────────────

def build_demo_pool_meta(pool_name: str) -> BalancerPoolMeta:
    """
    为演示直接构建已知池子的 PoolMeta（无需事件扫描）。
    生产环境中这些信息来自 scan_all_pools 或事件监听。
    """
    if pool_name == "80BAL-20WETH":
        pool_id = KNOWN_POOLS["80BAL-20WETH"]
        return BalancerPoolMeta(
            pool_id=pool_id,
            protocol="BalancerV2Weighted",
            tokens=[KNOWN_TOKENS["BAL"], KNOWN_TOKENS["WETH"]],
            token_decimals=[18, 18],
            pool_address=pool_id_to_address(pool_id),
            normalized_weights=[int(0.8 * ONE), int(0.2 * ONE)],
            swap_fee_percentage=int(0.0025 * ONE),
        )
    elif pool_name == "50WBTC-50WETH":
        pool_id = KNOWN_POOLS["50WBTC-50WETH"]
        return BalancerPoolMeta(
            pool_id=pool_id,
            protocol="BalancerV2Weighted",
            tokens=[KNOWN_TOKENS["WBTC"], KNOWN_TOKENS["WETH"]],
            token_decimals=[8, 18],
            pool_address=pool_id_to_address(pool_id),
            normalized_weights=[int(0.5 * ONE), int(0.5 * ONE)],
            swap_fee_percentage=int(0.003 * ONE),
        )
    elif pool_name == "50USDC-50WETH":
        pool_id = KNOWN_POOLS["50USDC-50WETH"]
        return BalancerPoolMeta(
            pool_id=pool_id,
            protocol="BalancerV2Weighted",
            tokens=[KNOWN_TOKENS["USDC"], KNOWN_TOKENS["WETH"]],
            token_decimals=[6, 18],
            pool_address=pool_id_to_address(pool_id),
            normalized_weights=[int(0.5 * ONE), int(0.5 * ONE)],
            swap_fee_percentage=int(0.003 * ONE),
        )
    else:
        raise ValueError(f"未知 pool 名称: {pool_name}，可选: {list(KNOWN_POOLS.keys())}")


def run_balancer_demo(w3: Web3, pool_name: str):
    """对指定 Balancer pool 运行完整的 6 步验证流程。"""
    print(f"\n{'='*65}")
    print(f" Balancer V2 协议解析验证：{pool_name}")
    print(f"{'='*65}\n")

    print("[ Step 1/4 ] 构建 PoolMeta")
    pool_meta = build_demo_pool_meta(pool_name)
    pool_meta_registry[normalize_pool_id(pool_meta.pool_id)] = pool_meta
    print(f"  pool_id:   {pool_meta.pool_id}")
    print(f"  address:   {pool_meta.pool_address}")
    print(f"  {pool_meta.summary()}\n")

    block = w3.eth.block_number
    print(f"  当前区块: {block:,}\n")

    print("[ Step 2/4 ] 获取链上 RawState")
    raw_state = fetch_raw_state(pool_meta, w3, block)
    if raw_state is None:
        print("  ❌ 获取 RawState 失败")
        return
    for i, (token, balance, weight) in enumerate(
        zip(raw_state.tokens, raw_state.balances, raw_state.normalized_weights)
    ):
        print(
            f"  token[{i}]: {token[:12]}...  "
            f"balance={balance / 10**pool_meta.token_decimals[i]:>18.4f}  "
            f"weight={weight/ONE:.0%}"
        )
    print(f"  fee:      {raw_state.swap_fee_percentage/ONE:.4%}\n")

    print("[ Step 3/4 ] RawState → Constraint")
    constraint = validate_and_convert(raw_state, pool_meta)
    if constraint is None:
        print("  ❌ Constraint 验证失败")
        return
    inv = constraint.calculate_invariant()
    print(f"  ✓ 验证通过  不变量 V = {inv:.6e}\n")

    print("[ Step 4/4 ] 精度验证（本地 vs 链上）")
    all_pass = True
    for i in range(len(constraint.balances)):
        for j in range(len(constraint.balances)):
            if i != j:
                passed = validate_constraint_accuracy(
                    constraint, w3, i, j, block, num_test_points=6, tolerance=0.0005
                )
                all_pass = all_pass and passed

    print("[ 最终状态 ]")
    status = check_pool_support_status(pool_meta.pool_id, w3, block)
    for k, v in status.to_dict().items():
        icon = "✓" if v is True else ("—" if v is None else ("✗" if v is False else " "))
        print(f"  {icon} {k}: {v}")


# ── Uniswap V3 演示 ──────────────────────────────────────

def run_v3_demo(w3: Web3, pool_name: str):
    """对指定 V3 pool 运行验证流程。"""
    print(f"\n{'='*65}")
    print(f" Uniswap V3 协议解析验证：{pool_name}")
    print(f"{'='*65}\n")

    if pool_name not in KNOWN_V3_POOLS:
        print(f"  ❌ 未知 V3 pool: {pool_name}")
        print(f"  可选: {list(KNOWN_V3_POOLS.keys())}")
        return

    info = KNOWN_V3_POOLS[pool_name]
    pool_meta = UniswapV3PoolMeta(
        pool_id=info["pool"].lower(),
        protocol="UniswapV3",
        tokens=[info["token0"], info["token1"]],
        token_decimals=info["decimals"],
        pool_address=info["pool"],
        fee=info["fee"],
        tick_spacing=FEE_TIER_TICK_SPACING[info["fee"]],
    )
    print(f"  {pool_meta.summary()}")
    print(f"  address: {pool_meta.pool_address}\n")

    adapter = UniswapV3Adapter()
    result = adapter.run_full_validation(pool_meta, w3)

    for step, passed in result.get("steps", {}).items():
        icon = "✓" if passed else "✗"
        print(f"  {icon} {step}: {passed}")

    for report in result.get("accuracy_reports", []):
        icon = "✓" if report["passed"] else "✗"
        print(f"  {icon} {report['direction']}: max_err={report['max_error']:.8%}")


# ── 扫描演示 ──────────────────────────────────────────────

def run_scan_demo(w3: Web3, lookback_blocks: int = 2000):
    """演示全量扫描。"""
    latest = w3.eth.block_number
    from_block = latest - lookback_blocks
    print(f"\n扫描最近 {lookback_blocks:,} 个区块 ({from_block:,} ~ {latest:,})")
    pools = scan_all_pools(w3, from_block, latest, batch_size=500)
    if pools:
        print(f"\n发现的 WeightedPool 样本（最多显示 3 个）：")
        for p in pools[:3]:
            print(f"  {p.pool_id[:20]}...  {p.summary()}")


def run_event_scan(args) -> int:
    """执行 protocol 事件扫描并产出报告文件。"""
    event_types = None
    if args.event_types.strip():
        event_types = [p.strip() for p in args.event_types.split(",") if p.strip()]

    result = scan_events_last_n_days(
        protocol=args.protocol,
        rpc=args.rpc,
        days=args.days,
        event_types=event_types,
        batch_size=args.batch_size,
        algebra_rpc=args.algebra_rpc,
        algebra_chain=args.algebra_chain,
        algebra_factory=args.algebra_factory,
    )
    weekly_report = build_weekly_report(result)
    paths = write_scan_outputs(
        scan_result=result,
        weekly_report=weekly_report,
        out_dir=args.out_dir,
    )

    print("\n事件扫描完成:")
    print(f"  run_id: {paths['run_id']}")
    print(f"  events: {paths['events_ndjson']}")
    print(f"  summary: {paths['summary_json']}")
    print(f"  analysis: {paths['analysis_md']}")

    totals = weekly_report.get("totals", {})
    print("\n统计摘要:")
    print(f"  total events: {totals.get('events_total', 0)}")
    print(f"  creation events: {totals.get('creation_events_total', 0)}")
    print(f"  state events: {totals.get('state_events_total', 0)}")
    print(f"  parse errors: {totals.get('parsed_error_total', 0)}")

    scan_meta = weekly_report.get("scan", {})
    warnings = scan_meta.get("warnings", [])
    if warnings:
        print(f"\nWarnings ({len(warnings)}):")
        for w in warnings[:10]:
            print(f"  - {w}")
    return 0


# ── 入口 ─────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="DexDecoder — DEX 协议解析框架")
    parser.add_argument(
        "--rpc",
        default=os.environ.get("ETH_RPC_URL", ""),
        help="以太坊 RPC URL（或设置环境变量 ETH_RPC_URL）"
    )
    parser.add_argument(
        "--protocol",
        default="balancer",
        choices=["all", "balancer", "v3", "v4", "algebra"],
        help="要验证的协议"
    )
    parser.add_argument(
        "--pool",
        default="",
        help="池子名称（协议相关）"
    )
    parser.add_argument(
        "--scan",
        action="store_true",
        help="扫描模式"
    )
    parser.add_argument(
        "--list-protocols",
        action="store_true",
        help="列出所有已注册协议"
    )
    parser.add_argument(
        "--event-scan",
        action="store_true",
        help="执行 protocol 事件扫描与分析"
    )
    parser.add_argument(
        "--days",
        type=int,
        default=7,
        help="事件扫描窗口天数（默认 7）"
    )
    parser.add_argument(
        "--out-dir",
        default="reports/events",
        help="事件扫描输出目录（默认 reports/events）"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=2000,
        help="日志扫描批大小（默认 2000）"
    )
    parser.add_argument(
        "--event-types",
        default="",
        help="可选：事件过滤（逗号分隔，支持 creation/state/具体事件名）"
    )
    parser.add_argument(
        "--algebra-rpc",
        default=os.environ.get("ALGEBRA_RPC_URL", ""),
        help="Algebra 扫描 RPC（可选；ethereum 模式默认复用 --rpc）"
    )
    parser.add_argument(
        "--algebra-chain",
        default="ethereum",
        choices=["ethereum", "polygon", "arbitrum"],
        help="Algebra 部署链（默认 ethereum）"
    )
    parser.add_argument(
        "--algebra-factory",
        default=os.environ.get("ALGEBRA_FACTORY_ADDRESS", ""),
        help="Algebra Factory 地址（ethereum 模式必填）"
    )
    args = parser.parse_args()

    if args.list_protocols:
        registry = build_registry()
        print(registry.summary())
        return

    if args.event_scan:
        need_eth_rpc = args.protocol in {"all", "balancer", "v3", "v4"}
        need_algebra_rpc = args.protocol in {"all", "algebra"}
        if need_eth_rpc and not args.rpc:
            print("❌ event scan 需要以太坊 RPC：--rpc 或 ETH_RPC_URL")
            return
        if need_algebra_rpc and not (args.algebra_rpc or args.rpc):
            print("❌ event scan 的 algebra 需要 --algebra-rpc 或 --rpc")
            return
        if need_algebra_rpc and args.algebra_chain == "ethereum" and not args.algebra_factory:
            print("❌ algebra_chain=ethereum 时必须提供 --algebra-factory（或 ALGEBRA_FACTORY_ADDRESS）")
            return
        if args.days <= 0:
            print("❌ --days 必须大于 0")
            return
        if args.batch_size <= 0:
            print("❌ --batch-size 必须大于 0")
            return
        exit_code = run_event_scan(args)
        if exit_code != 0:
            sys.exit(exit_code)
        return

    if not args.rpc:
        print("❌ 请提供 RPC URL：--rpc https://... 或设置 ETH_RPC_URL 环境变量")
        sys.exit(1)

    print(f"连接 RPC: {args.rpc[:40]}...")
    w3 = Web3(Web3.HTTPProvider(args.rpc))
    if not w3.is_connected():
        print("❌ RPC 连接失败")
        sys.exit(1)

    chain_id = w3.eth.chain_id
    latest_block = w3.eth.block_number
    print(f"✓ 已连接  chain_id={chain_id}  latest_block={latest_block:,}\n")

    if args.scan:
        run_scan_demo(w3)
    elif args.protocol == "balancer":
        pool = args.pool or "80BAL-20WETH"
        run_balancer_demo(w3, pool)
    elif args.protocol == "v3":
        pool = args.pool or "USDC-WETH-3000"
        run_v3_demo(w3, pool)
    else:
        print(f"协议 '{args.protocol}' 的完整演示脚本尚未添加。")
        print("请使用 Python API 直接调用 Adapter：")
        print(f"  from src import {args.protocol.capitalize()}Adapter")


if __name__ == "__main__":
    main()