#!/usr/bin/env python3
"""
Raydium MEV Trader - Quantum Beast Edition
-----------------------------------------
A high-performance MEV trading bot that:
1. Monitors Raydium 0x03 swap instructions via WebSocket
2. Verifies pool security (mint authority revoked, liquidity locked)
3. Front-runs trades with exact route mirroring
4. Optimized for minimal latency with Rust backend

Created by QuantumBeast - Enhanced by Claude
"""

import os
import sys
import json
import time
import base64
import base58
import asyncio
import platform
import asyncio
from asyncio import Semaphore
import websockets
import traceback
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional, Set, Union, Tuple

# Import Rust trading core
try:
    from solana_rust_bot import RaydiumTradingCore, SwapParams, TokenPairSetting, TradingConfig
except ImportError:
    print("❌ ERROR: Rust trading core not found. Make sure to build the Rust extension first.")
    print("Run: maturin develop --release")
    sys.exit(1)

try:
    from solana_logger import logger
    from config_loader import config, load_env_file
except ImportError:
    print("❌ ERROR: Custom modules not found. Check your project structure.")
    sys.exit(1)

# Constants
LAMPORTS_PER_SOL = 1_000_000_000
MIN_TRADE_INTERVAL = float(config.get("min_trade_interval", 1.0))  # Increased from 0.2 to 1.0
MAX_ACTIVE_MEVS = int(config.get("max_active_mevs", 1))
FRONT_RUN_AMOUNT = float(config.get("front_run_amount", 0.3))  # Updated to use 0.3 as default
MIN_VICTIM_SIZE = float(config.get("sandwich_min_victim_size", 10.0))  # Increased from 5.0 to 10.0
MAX_VICTIM_SIZE = float(config.get("sandwich_max_victim_size", 320.0))
MIN_PROFIT = float(config.get("min_profit", 0.002))
MIN_PROFIT_PCT = float(config.get("min_profit_pct", 0.015))  # Minimum 1.5% profit
MAX_RETRIES_ON_ERROR = 3
ERROR_COOLDOWN = 30  # seconds to wait after a critical error

# State tracking
last_trade_time = 0
active_trades = 0
seen_txs = set()
pool_security_cache = {}
# Fix for Windows event loop
if platform.system() == "Windows":
    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

# Rate limiter class
class RateLimiter:
    """Simple rate limiter to prevent hitting API limits"""
    
    def __init__(self, max_calls_per_second=2):  # Reduced from 5 to 2
        self.max_calls = max_calls_per_second
        self.semaphore = Semaphore(max_calls_per_second)
        self.call_times = []
    
    async def acquire(self):
        """Acquire permission to make a call, respecting the rate limit"""
        # Clean up old timestamps
        current_time = time.time()
        self.call_times = [t for t in self.call_times if current_time - t < 1.0]
        
        # If we've hit our limit, wait until we're under it
        if len(self.call_times) >= self.max_calls:
            wait_time = 1.0 - (current_time - self.call_times[0])
            if wait_time > 0:
                await asyncio.sleep(wait_time)
        
        # Acquire semaphore
        await self.semaphore.acquire()
        
        # Record this call time
        self.call_times.append(time.time())
        
    def release(self):
        """Release the semaphore after a call is complete"""
        self.semaphore.release()

class MEVTrader:
    """High-performance MEV trading bot for Raydium"""
    
    def __init__(self):
        """Initialize the MEV trader"""
        self.helius_ws_url = config.get("helius_ws_url")
        self.raydium_program_id = config.get("raydium_program_id")
        
        # Load RPC URLs for rotation
        self.rpc_urls = config.get("rpc_urls", [config.get("rpc_url")])
        if isinstance(self.rpc_urls, str):
            self.rpc_urls = [self.rpc_urls]
        self.current_rpc_index = 0
        
        # Load trading core from Rust
        private_key_path = config.get("keypair_path")
        if not os.path.exists(private_key_path):
            logger.error(f"Keypair file not found: {private_key_path}")
            sys.exit(1)
            
        logger.info(f"Initializing Raydium Trading Core with keypair: {private_key_path}")
        self.trading_core = RaydiumTradingCore(
            rpc_url=self.get_next_rpc_url(),
            keypair_path=private_key_path
        )
        
        # System state
        self.connected = False
        self.subscription_id = None
        self.last_trade_time = 0
        self.active_trades = 0
        self.seen_txs = set()
        self.trading_enabled = config.get_bool("trading", True)
        self.pool_security_cache = {}
        self.trades_executed = 0
        self.profits = 0.0
        self.total_trade_volume = 0.0
        self.successful_trades = 0
        self.failed_trades = 0
        self.start_time = time.time()
        self.error_count = 0
        self.last_error_time = 0
        
        # Trade history
        self.trade_history = []
        self.max_history_items = 1000
        
        # Initialize balances
        self.sol_balance = 0.0
        self.wsol_balance = 0.0
        self.initial_sol_balance = 0.0
        self.initial_wsol_balance = 0.0
        
        # For profit calculation
        self.session_gas_spent = 0.0
        
        # Initialize rate limiter
        self.rate_limiter = RateLimiter(max_calls_per_second=2)  # Reduced from 4 to 2
        
        # Set up auto-restart
        self.auto_restart = config.get_bool("auto_restart", True)
        self.restart_attempts = 0
        self.max_restart_attempts = 10
        
        # Create data directory for logs and trade history
        self.data_dir = Path("data")
        self.data_dir.mkdir(exist_ok=True)
        
        logger.info(f"MEV Trader initialized with public key: {self.trading_core.get_public_key()}")
        logger.info(f"Trading enabled: {self.trading_enabled}")
        logger.info(f"Using {len(self.rpc_urls)} RPC endpoints for load balancing")
    
    def get_next_rpc_url(self):
        """Get the next RPC URL in the rotation"""
        url = self.rpc_urls[self.current_rpc_index]
        self.current_rpc_index = (self.current_rpc_index + 1) % len(self.rpc_urls)
        return url
        
    async def get_balances(self):
        """Get wallet balances using RPC call (only used once at startup)"""
        try:
            # Apply rate limiting before making RPC call
            await self.rate_limiter.acquire()
            try:
                # Add artificial delay before RPC call
                await asyncio.sleep(0.5)
                
                sol_balance, token_balances = await self.trading_core.get_balances()
                self.sol_balance = sol_balance
                self.wsol_balance = token_balances.get(config.get("sol_mint"), 0.0)
                
                # Store initial balances for profit calculation
                if self.initial_sol_balance == 0.0:
                    self.initial_sol_balance = sol_balance
                if self.initial_wsol_balance == 0.0:
                    self.initial_wsol_balance = self.wsol_balance
                
                logger.info(f"Initial wallet balances - SOL: {self.sol_balance:.6f}, WSOL: {self.wsol_balance:.6f}")
                
                # Verify we have enough balance
                min_sol = float(config.get("min_sol_balance", 0.03))
                min_wsol = float(config.get("min_wsol_balance", 0.02))
                
                if self.sol_balance < min_sol:
                    logger.warning(f"SOL balance too low: {self.sol_balance:.6f} < {min_sol:.6f}")
                
                if self.wsol_balance < min_wsol:
                    logger.warning(f"WSOL balance too low: {self.wsol_balance:.6f} < {min_wsol:.6f}")
                    
                return True
            finally:
                # Release rate limiter after RPC call completes
                self.rate_limiter.release()
        except Exception as e:
            logger.error(f"Error getting balances: {e}")
            return False
            
    async def start_balance_monitoring(self):
        """Start WebSocket monitoring of SOL and WSOL balances"""
        try:
            # Get account addresses from config
            sol_pubkey = config.get("solana_public_key")
            wsol_token_account = config.get("wsol_token_account")
            
            if not sol_pubkey or not wsol_token_account:
                logger.error("Missing solana_public_key or wsol_token_account in config")
                return False
                
            logger.info(f"Starting WebSocket balance monitoring")
            logger.info(f"SOL account: {sol_pubkey}")
            logger.info(f"WSOL account: {wsol_token_account}")
            
            # Use API key if available
            headers = {}
            api_key = config.get("api_key")
            if api_key:
                headers["api-key"] = api_key
            
            async with websockets.connect(
                self.helius_ws_url,
                extra_headers=headers,
                ping_interval=20,
                ping_timeout=30,
                close_timeout=60
            ) as ws:
                logger.info("✅ Connected to WebSocket for balance monitoring")
                
                # Add delay before subscription to avoid rate limiting
                await asyncio.sleep(1.0)
                
                # Subscribe to SOL account
                await ws.send(
                    json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "accountSubscribe",
                        "params": [
                            sol_pubkey,
                            {"encoding": "base64", "commitment": "confirmed"}
                        ]
                    })
                )
                
                # Add delay between subscriptions
                await asyncio.sleep(1.0)
                
                # Subscribe to WSOL token account
                await ws.send(
                    json.dumps({
                        "jsonrpc": "2.0",
                        "id": 2,
                        "method": "accountSubscribe",
                        "params": [
                            wsol_token_account,
                            {"encoding": "base64", "commitment": "confirmed"}
                        ]
                    })
                )
                
                # Track subscriptions
                subs = {}
                
                # Process subscription confirmations
                for _ in range(2):  # We expect 2 confirmation messages
                    msg = await ws.recv()
                    data = json.loads(msg)
                    
                    if "result" in data and "id" in data:
                        if data["id"] == 1:
                            subs[data["result"]] = "SOL"
                            logger.info(f"SOL balance subscription confirmed (ID: {data['result']})")
                        elif data["id"] == 2:
                            subs[data["result"]] = "WSOL"
                            logger.info(f"WSOL balance subscription confirmed (ID: {data['result']})")
                
                # Monitor balance updates
                while True:
                    try:
                        msg = await asyncio.wait_for(ws.recv(), timeout=60)
                        data = json.loads(msg)
                        
                        if "method" in data and data["method"] == "accountNotification":
                            sub_id = data["params"]["subscription"]
                            acc_type = subs.get(sub_id, "Unknown")
                            
                            if acc_type == "SOL":
                                prev_balance = self.sol_balance
                                lamports = data["params"]["result"]["value"]["lamports"]
                                self.sol_balance = lamports / LAMPORTS_PER_SOL
                                logger.info(f"SOL Balance updated: {self.sol_balance:.6f} ({self.sol_balance - prev_balance:+.6f})")
                                
                                # Check if this was a trade completion
                                if abs(self.sol_balance - prev_balance) > 0.000001:
                                    # Update profit calculations with new balance
                                    self.update_balance_and_profit()
                                    
                            elif acc_type == "WSOL":
                                prev_balance = self.wsol_balance
                                account_data = data["params"]["result"]["value"]["data"][0]
                                self.wsol_balance = self._extract_token_balance(account_data)
                                logger.info(f"WSOL Balance updated: {self.wsol_balance:.6f} ({self.wsol_balance - prev_balance:+.6f})")
                                
                                # Check if this was a trade completion
                                if abs(self.wsol_balance - prev_balance) > 0.000001:
                                    # Update profit calculations with new balance
                                    self.update_balance_and_profit()
                                
                    except asyncio.TimeoutError:
                        # Send a ping to keep connection alive
                        await ws.ping()
                    except websockets.ConnectionClosed:
                        logger.warning("WebSocket connection closed. Reconnecting...")
                        # Add delay before reconnection attempt
                        await asyncio.sleep(2.0)
                        return await self.start_balance_monitoring()
        
        except Exception as e:
            logger.error(f"Error in balance monitoring: {e}")
            logger.error(traceback.format_exc())
            await asyncio.sleep(5)  # Wait before reconnecting
            return await self.start_balance_monitoring()
    
    def _extract_token_balance(self, account_data_b64: str) -> float:
        """Extract token balance from base64 encoded SPL token account data"""
        try:
            # Decode the base64 string
            account_data = base64.b64decode(account_data_b64)
            
            # SPL Token account layout:
            # - offset 0-31: mint address
            # - offset 32-63: owner address
            # - offset 64-71: amount (u64, little endian)
            # - remaining bytes: other fields
            
            # Token amount is at offset 64 (8 bytes)
            if len(account_data) >= 72:
                amount_bytes = account_data[64:72]
                amount = int.from_bytes(amount_bytes, byteorder="little")
                return amount / LAMPORTS_PER_SOL
            return 0.0
        except Exception as e:
            logger.error(f"Error extracting token balance: {e}")
            return 0.0
            
    def add_trade_to_history(self, trade_record: Dict):
        """Add trade to history and trim if needed"""
        # Add to history
        self.trade_history.append(trade_record)
        
        # Trim if needed
        if len(self.trade_history) > self.max_history_items:
            self.trade_history = self.trade_history[-self.max_history_items:]
        
        # Save to disk (async in production)
        try:
            # Save to JSON file
            trade_file = self.data_dir / "trade_history.json"
            with open(trade_file, "w") as f:
                json.dump(self.trade_history, f)
        except Exception as e:
            logger.error(f"Error saving trade history: {e}")
    
    def calculate_profit_stats(self):
        """Calculate profit statistics"""
        total_profit = 0.0
        profit_pct = 0.0
        
        # Calculate from current balances
        current_total = self.sol_balance + self.wsol_balance
        initial_total = self.initial_sol_balance + self.initial_wsol_balance
        
        if initial_total > 0:
            total_profit = current_total - initial_total
            profit_pct = (total_profit / initial_total) * 100
        
        # Subtract gas costs (simplified)
        total_profit -= self.session_gas_spent
        
        # Calculate win rate
        win_rate = 0
        if self.trades_executed > 0:
            win_rate = (self.successful_trades / self.trades_executed) * 100
            
        return {
            "total_profit_sol": total_profit,
            "profit_percentage": profit_pct,
            "trades_executed": self.trades_executed,
            "successful_trades": self.successful_trades,
            "failed_trades": self.failed_trades,
            "win_rate": win_rate,
            "gas_spent": self.session_gas_spent,
            "uptime_seconds": time.time() - self.start_time
        }
        
    def log_profit_stats(self):
        """Log profit statistics"""
        stats = self.calculate_profit_stats()
        
        # Format uptime
        uptime_seconds = stats["uptime_seconds"]
        hours, remainder = divmod(uptime_seconds, 3600)
        minutes, seconds = divmod(remainder, 60)
        uptime_str = f"{int(hours)}h {int(minutes)}m {int(seconds)}s"
        
        # Log stats
        logger.info("\n" + "="*80)
        logger.info("📊 PROFIT STATISTICS")
        logger.info("="*80)
        logger.info(f"Total Profit: {stats['total_profit_sol']:.6f} SOL ({stats['profit_percentage']:.2f}%)")
        logger.info(f"Trades Executed: {stats['trades_executed']} (Success: {stats['successful_trades']}, Failed: {stats['failed_trades']})")
        logger.info(f"Win Rate: {stats['win_rate']:.2f}%")
        logger.info(f"Gas Spent: {stats['gas_spent']:.6f} SOL")
        logger.info(f"Uptime: {uptime_str}")
        logger.info("="*80 + "\n")
        
        # Save stats to file
        try:
            stats_file = self.data_dir / "profit_stats.json"
            with open(stats_file, "w") as f:
                json.dump(stats, f)
        except Exception as e:
            logger.error(f"Error saving profit stats: {e}")
    
    def update_balance_and_profit(self):
        """Update profit calculations after balance updates"""
        # Calculate profit from trades in history that haven't been processed
        for trade in self.trade_history:
            if trade["status"] == "success" and "profit_calculated" not in trade:
                # Mark as processed
                trade["profit_calculated"] = True
                
                # This is where you would do exact profit calculation
                # For now, we're using the estimated profit already in the trade record
                if "estimated_profit_sol" in trade:
                    self.profits += trade["estimated_profit_sol"]
                    
        # Log updated profit stats
        self.log_profit_stats()
    async def verify_pool_security(self, pool_address: str) -> Tuple[bool, bool, bool]:
        """
        Verify pool security by checking:
        1. Mint authority is revoked
        2. Liquidity is locked
        3. Has sufficient exit liquidity
        
        Args:
            pool_address: The pool account address
            
        Returns:
            (is_secure, has_exit_liquidity, in_cache)
        """
        try:
            # Check cache first
            if pool_address in self.pool_security_cache:
                cache_entry = self.pool_security_cache[pool_address]
                # Return cached result if less than 10 minutes old
                if time.time() - cache_entry["timestamp"] < 600:
                    return cache_entry["is_secure"], cache_entry["has_exit_liquidity"], True
            
            # Fetch pool account data
            logger.debug(f"Fetching pool data for {pool_address}")
            
            # Apply rate limiting before making RPC call
            await self.rate_limiter.acquire()
            try:
                # Add artificial delay before RPC call
                await asyncio.sleep(1.0)
                
                # Build RPC request
                request = {
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "getAccountInfo",
                    "params": [
                        pool_address,
                        {"encoding": "base64", "commitment": "confirmed"}
                    ]
                }
                
                # Use a different endpoint for each request to distribute load
                ws_url = self.get_next_rpc_url()
                async with websockets.connect(ws_url) as ws:
                    await ws.send(json.dumps(request))
                    response = await ws.recv()
                    data = json.loads(response)
                    
                    if "result" not in data or data["result"] is None or data["result"]["value"] is None:
                        logger.warning(f"Pool {pool_address} not found")
                        return False, False, False
                    
                    account_data_b64 = data["result"]["value"]["data"][0]
                    account_data = base64.b64decode(account_data_b64)
                    
                    # Check discriminator (first byte should be 0x06)
                    if len(account_data) < 1 or account_data[0] != 0x06:
                        logger.warning(f"Pool {pool_address} is not a valid Raydium pool")
                        return False, False, False
                    
                    # Check mint authority revoked (offset 72, value should be 0xF4)
                    mint_revoked = False
                    if len(account_data) > 72:
                        mint_byte = account_data[72]
                        mint_revoked = (mint_byte == 0xF4)
                    
                    # Check liquidity locked (offset 88, should be non-zero)
                    lp_locked = False
                    if len(account_data) >= 96:  # 88 + 8 bytes
                        lp_raw_value = int.from_bytes(account_data[88:96], "little")
                        lp_locked = (lp_raw_value > 0)
                    
                    # Get token info and reserves to check exit liquidity
                    has_exit_liquidity = False
                    sol_reserve = 0.0
                    
                    # Extract base and quote mints (offsets from Raydium v4 layout)
                    if len(account_data) >= 432:  # 400 + 32
                        base_mint_data = account_data[368:400]
                        quote_mint_data = account_data[400:432]
                        
                        base_mint = base58.b58encode(base_mint_data).decode('utf-8')
                        quote_mint = base58.b58encode(quote_mint_data).decode('utf-8')
                        
                        # Extract reserves (offsets from Raydium v4 layout)
                        if len(account_data) >= 296:  # 288 + 8
                            base_reserve = int.from_bytes(account_data[256:264], "little")
                            quote_reserve = int.from_bytes(account_data[288:296], "little")
                            
                            # Check which one is SOL
                            sol_mint = "So11111111111111111111111111111111111111112"
                            if base_mint == sol_mint:
                                sol_reserve = base_reserve / LAMPORTS_PER_SOL
                            elif quote_mint == sol_mint:
                                sol_reserve = quote_reserve / LAMPORTS_PER_SOL
                            
                            # Check exit liquidity
                            min_exit_liquidity = float(config.get("min_exit_liquidity_sol", 5.0))
                            has_exit_liquidity = sol_reserve >= min_exit_liquidity
                    
                    # Cache the result
                    is_secure = mint_revoked and lp_locked
                    self.pool_security_cache[pool_address] = {
                        "is_secure": is_secure,
                        "has_exit_liquidity": has_exit_liquidity,
                        "timestamp": time.time()
                    }
                    
                    # Log detailed security info
                    security_status = "SECURE ✅" if is_secure else "INSECURE ❌"
                    liquidity_status = "SUFFICIENT ✅" if has_exit_liquidity else "INSUFFICIENT ❌"
                    logger.info(f"Pool {pool_address}: {security_status} | Exit Liquidity: {liquidity_status} ({sol_reserve:.2f} SOL)")
                    
                    return is_secure, has_exit_liquidity, False
            finally:
                # Release rate limiter after RPC call completes
                self.rate_limiter.release()
        
        except Exception as e:
            logger.error(f"Error verifying pool {pool_address}: {e}")
            logger.error(traceback.format_exc())
            return False, False, False
        
    def decode_ray_log(self, log_line: str) -> Optional[Dict]:
        """
        Decode Raydium ray_log data from log line
        Returns decoded data dictionary or None
        """
        try:
            if "ray_log:" not in log_line:
                return None
                
            # Extract the base64 encoded part after ray_log:
            ray_log_part = log_line.split("ray_log:")[1].strip()
            raw_data = base64.b64decode(ray_log_part)
            
            if len(raw_data) < 1:
                return None
                
            # Get instruction type (first byte)
            instr_type = raw_data[0]
            instr_hex = f"0x{instr_type:02x}"
            
            # Check if it's a 0x03 instruction (swaps)
            if instr_type == 0x03:
                # Check if we have enough data for amount
                if len(raw_data) >= 33:
                    # Extract values (uint64, 8 bytes each)
                    values = {}
                    values["instruction_type"] = instr_hex
                    
                    # Extract values
                    for i in range(4):
                        start_pos = 1 + (i * 8)
                        end_pos = start_pos + 8
                        
                        if len(raw_data) >= end_pos:
                            val = int.from_bytes(raw_data[start_pos:end_pos], byteorder="little")
                            values[f"value_{i+1}"] = val
                            values[f"value_{i+1}_sol"] = val / LAMPORTS_PER_SOL
                    
                    # Add raw data
                    values["raw_data"] = raw_data
                    values["raw_data_hex"] = raw_data.hex()
                    values["timestamp"] = datetime.now().isoformat()
                    
                    return values
            
            return None
                
        except Exception as e:
            logger.error(f"Error decoding ray_log: {e}")
            return None
    
    async def process_victim_swap(self, tx_data: Dict, signature: str) -> bool:
        """Process a victim's swap and execute a MEV opportunity"""
        global active_trades, last_trade_time
        
        # Check trade throttling
        if active_trades >= MAX_ACTIVE_MEVS:
            logger.debug(f"Max active MEVs reached: {active_trades}/{MAX_ACTIVE_MEVS}")
            return False
            
        current_time = time.time()
        if current_time - last_trade_time < MIN_TRADE_INTERVAL:
            logger.debug(f"Trade throttled: {current_time - last_trade_time:.3f}s < {MIN_TRADE_INTERVAL:.3f}s")
            return False
        
        # Get victim swap amount
        victim_amount_sol = tx_data.get("value_1_sol", 0)
        
        # Check victim size constraints
        if victim_amount_sol < MIN_VICTIM_SIZE:
            logger.debug(f"Victim size too small: {victim_amount_sol:.6f} SOL < {MIN_VICTIM_SIZE:.1f} SOL")
            return False
            
        if victim_amount_sol > MAX_VICTIM_SIZE:
            logger.debug(f"Victim size too large: {victim_amount_sol:.6f} SOL > {MAX_VICTIM_SIZE:.1f} SOL")
            return False
        
        # Check WSOL balance
        if self.wsol_balance < FRONT_RUN_AMOUNT:
            logger.error(f"Insufficient WSOL balance: {self.wsol_balance:.6f} < {FRONT_RUN_AMOUNT:.2f}")
            return False
        
        # Set trading in progress flags
        active_trades += 1
        last_trade_time = current_time
        
        # Create trade record for history
        trade_record = {
            "timestamp": datetime.now().isoformat(),
            "victim_signature": signature,
            "victim_amount_sol": victim_amount_sol,
            "front_run_amount_sol": FRONT_RUN_AMOUNT,
            "status": "pending",
            "profit_sol": 0.0,
            "profit_percent": 0.0,
            "gas_used": 0.0
        }
        
        try:
            # Extract instruction data for debugging
            instruction_data = tx_data["raw_data"]
            logger.debug(f"Original instruction data: {instruction_data.hex()}")
            
            # Modify instruction data to use our front-run amount
            front_run_lamports = int(FRONT_RUN_AMOUNT * LAMPORTS_PER_SOL)
            new_instruction_data = bytearray(instruction_data)
            
            # Replace amount (bytes 1-9) with our amount in little-endian
            amount_bytes = front_run_lamports.to_bytes(8, byteorder="little")
            for i in range(8):
                if i + 1 < len(new_instruction_data):
                    new_instruction_data[i+1] = amount_bytes[i]
            
            # *** CRITICAL FIX: Change instruction type from 0x03 to 0x06 ***
            # Raydium's API uses 0x06 as the swap instruction, while mempool shows 0x03
            if new_instruction_data[0] == 0x03:
                new_instruction_data[0] = 0x06
                logger.debug(f"Changed instruction type from 0x03 to 0x06 for Raydium compatibility")
                
            logger.debug(f"Modified instruction data: {bytes(new_instruction_data).hex()}")
            
            # Add artificial delay before RPC call to reduce rate limiting
            logger.debug(f"Adding artificial delay before RPC call")
            await asyncio.sleep(1.0)
            
            logger.trade(f"Executing front-run with {FRONT_RUN_AMOUNT:.3f} SOL")
            
            # Track balances before trade for profit calculation
            pre_sol_balance = self.sol_balance
            pre_wsol_balance = self.wsol_balance
            
            # Apply rate limiting for RPC call
            await self.rate_limiter.acquire()
            try:
                # Execute the trade with exponential backoff for 429 errors
                if self.trading_enabled:
                    # For a true dynamic mirror trader, you should attempt to extract the actual token
                    # from transaction logs. For now we'll use a default value which your lib.rs will handle
                    token_mint = "DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263"  # BONK as default
                    
                    # Use exponential backoff for retries
                    retries = 3
                    backoff = 1.0
                    
                    while retries > 0:
                        try:
                            # Update the RPC URL to a different one for each retry
                            self.trading_core.update_rpc_url(self.get_next_rpc_url())
                            
                            # Timestamp for latency tracking
                            start_time = time.time()
                            
                            # Your lib.rs will dynamically create token pair settings if they don't exist
                            result = await self.trading_core.execute_custom_swap(token_mint, bytes(new_instruction_data))
                            
                            # Calculate latency
                            latency_ms = (time.time() - start_time) * 1000
                            
                            if result.success:
                                # Update trade metrics
                                self.trades_executed += 1
                                self.successful_trades += 1
                                
                                # Calculate profit (this will be updated properly when balance updates)
                                estimated_profit = victim_amount_sol * MIN_PROFIT_PCT
                                
                                # Update trade record
                                trade_record["status"] = "success"
                                trade_record["tx_signature"] = result.signature
                                trade_record["token_mint"] = token_mint
                                trade_record["estimated_profit_sol"] = estimated_profit
                                trade_record["latency_ms"] = latency_ms
                                
                                # Add to trade history
                                self.add_trade_to_history(trade_record)
                                
                                # Log the success
                                logger.profit(f"Front-run success! TX: {result.signature} | Est. Profit: ~{estimated_profit:.6f} SOL | Latency: {latency_ms:.2f}ms")
                                return True
                            else:
                                # Update metrics
                                self.failed_trades += 1
                                
                                # Update trade record
                                trade_record["status"] = "failed"
                                trade_record["error"] = result.error
                                
                                # Log the error
                                logger.error(f"Front-run failed: {result.error}")
                                
                                # If we got a rate limit error, retry with backoff
                                if "429 Too Many Requests" in result.error and retries > 1:
                                    logger.warning(f"Rate limited (429). Backing off for {backoff}s. Retries left: {retries-1}")
                                    await asyncio.sleep(backoff)
                                    backoff *= 2  # Exponential backoff
                                    retries -= 1
                                    continue
                                
                                # Add to trade history
                                self.add_trade_to_history(trade_record)
                                return False
                                
                        except ValueError as e:
                            # Handle rate limiting errors
                            if "429 Too Many Requests" in str(e) and retries > 1:
                                logger.warning(f"Rate limited (429). Backing off for {backoff}s. Retries left: {retries-1}")
                                await asyncio.sleep(backoff)
                                backoff *= 2  # Exponential backoff
                                retries -= 1
                            else:
                                # Update trade record
                                trade_record["status"] = "error"
                                trade_record["error"] = str(e)
                                
                                # Add to trade history
                                self.add_trade_to_history(trade_record)
                                
                                logger.error(f"Error executing front-run: {e}")
                                return False
                        
                    return False
                else:
                    logger.warning("Trading disabled. Simulation only.")
                    return True
            finally:
                # Release rate limiter
                self.rate_limiter.release()
                
        except Exception as e:
            # Update trade record
            trade_record["status"] = "error"
            trade_record["error"] = str(e)
            
            # Add to trade history
            self.add_trade_to_history(trade_record)
            
            logger.error(f"Error executing front-run: {e}")
            traceback.print_exc()
            return False
        finally:
            # Reset flag
            active_trades -= 1
    
   async def monitor_ray_logs(self):
    """Monitor Raydium logs specifically focusing on 0x03 instructions"""
    global active_trades  # Move global declaration to the beginning of the function
    
    retry_delay = 1
    max_retry_delay = 30
    
    logger.info("Starting Raydium MEV Monitor")
    print("\n" + "="*80)
    print("🚀 RAYDIUM MEV TRADER - QUANTUM BEAST EDITION")
    print("="*80 + "\n")
    print("Monitoring for profitable 0x03 swap opportunities...")
    print(f"Target: {MIN_VICTIM_SIZE}-{MAX_VICTIM_SIZE} SOL swaps")
    print(f"Front-run amount: {FRONT_RUN_AMOUNT} SOL")
    print(f"Rate limiting: Max {self.rate_limiter.max_calls} RPC calls per second")
    print(f"Trade interval: Minimum {MIN_TRADE_INTERVAL} seconds between trades")
    
    # Initialize balances first
    await self.get_balances()
    
    while True:
        try:
            # Connect to WebSocket with a random endpoint from our pool
            ws_url = self.get_next_rpc_url()
            if not ws_url.startswith("wss://"):
                # Convert RPC URL to WebSocket URL if needed
                ws_url = ws_url.replace("https://", "wss://")
            
            logger.info(f"Connecting to WebSocket: {ws_url}")
            
            # Add delay before connection to avoid rate limiting
            await asyncio.sleep(2.0)
            
            async with websockets.connect(
                self.helius_ws_url,
                ping_interval=20,
                ping_timeout=30,
                close_timeout=60
            ) as ws:
                logger.info("✅ Connected to WebSocket")
                self.connected = True
                
                # Add delay before subscription to avoid rate limiting
                await asyncio.sleep(1.0)
                
                # Subscribe to Raydium program logs
                subscription_msg = json.dumps({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "logsSubscribe",
                    "params": [
                        {"mentions": [self.raydium_program_id]},
                        {"commitment": "processed"}
                    ]
                })
                
                await ws.send(subscription_msg)
                
                # Check subscription confirmation
                response = await ws.recv()
                confirmation = json.loads(response)
                
                if "result" in confirmation:
                    self.subscription_id = confirmation["result"]
                    logger.info(f"✅ Subscription successful (ID: {self.subscription_id})")
                    print(f"Watching for 0x03 MEV opportunities...\n")
                    # Reset retry delay on successful connection
                    retry_delay = 1
                    
                    # Reset active trades counter on connection
                    active_trades = 0
                else:
                    logger.error(f"❌ Subscription failed: {confirmation}")
                    raise Exception("Failed to subscribe to logs")
                
                # Process WebSocket messages
                while True:
                    try:                            
                        # Receive message with timeout
                        msg = await asyncio.wait_for(ws.recv(), timeout=60)
                        
                        try:
                            data = json.loads(msg)
                        except json.JSONDecodeError:
                            continue
                        
                        # Check if valid notification
                        if "params" in data and "result" in data["params"]:
                            result = data["params"]["result"]
                            
                            # Extract logs and signature
                            if "value" in result:
                                value = result["value"]
                                logs = value.get("logs", [])
                                signature = value.get("signature", "unknown")
                                
                                # Skip if no logs or already seen
                                if not logs or signature in self.seen_txs:
                                    continue
                                    
                                # Add to seen transactions
                                self.seen_txs.add(signature)
                                if len(self.seen_txs) > 5000:
                                    self.seen_txs = set(list(self.seen_txs)[-2500:])
                                
                                # Check for ray_log entries
                                for log in logs:
                                    if "ray_log:" in log:
                                        # Decode the ray_log
                                        decoded = self.decode_ray_log(log)
                                        
                                        # If it's a 0x03 instruction, process it
                                        if decoded and decoded["instruction_type"] == "0x03":
                                            logger.info(f"🔎 Found 0x03 swap: {signature[:16]}...")
                                            logger.info(f"    Amount: {decoded['value_1_sol']:.6f} SOL")
                                            
                                            # Check if this is a tradable MEV opportunity
                                            if MIN_VICTIM_SIZE <= decoded["value_1_sol"] <= MAX_VICTIM_SIZE:
                                                logger.info(f"💰 Potential MEV opportunity: {decoded['value_1_sol']:.6f} SOL")
                                                
                                                # Add delay to spread out processing of opportunities
                                                await asyncio.sleep(0.5)
                                                
                                                # Process the opportunity in a separate task
                                                asyncio.create_task(self.process_victim_swap(decoded, signature))
                        
                    except asyncio.TimeoutError:
                        # Send a ping to keep connection alive
                        await ws.ping()
                    except websockets.ConnectionClosed:
                        logger.warning("WebSocket connection closed.")
                        raise
        
        except Exception as e:
            self.connected = False
            logger.error(f"Error: {e}")
            logger.info(f"Reconnecting in {retry_delay} seconds...")
            await asyncio.sleep(retry_delay)
            retry_delay = min(retry_delay * 2, max_retry_delay)
            
            # Reset active trades counter on reconnection
            active_trades = 0

    async def main(self):
        """Main entry point"""
        try:
            # Print startup banner
            logger.info("\n" + "="*80)
            logger.info("🚀 RAYDIUM MEV TRADER - QUANTUM BEAST EDITION")
            logger.info("="*80)
            logger.info(f"Bot Public Key: {self.trading_core.get_public_key()}")
            logger.info(f"Trading Enabled: {self.trading_enabled}")
            logger.info(f"Front-run Amount: {FRONT_RUN_AMOUNT} SOL")
            logger.info(f"Target Size Range: {MIN_VICTIM_SIZE}-{MAX_VICTIM_SIZE} SOL")
            logger.info(f"Minimum Profit: {MIN_PROFIT_PCT*100:.1f}%")
            logger.info(f"Rate Limiting: Max {self.rate_limiter.max_calls} RPC calls/second")
            logger.info(f"Trade Interval: Minimum {MIN_TRADE_INTERVAL} seconds between trades")
            
            # First get initial balances
            await self.get_balances()
            
            # Load trade history if exists
            trade_file = self.data_dir / "trade_history.json"
            if trade_file.exists():
                try:
                    with open(trade_file, "r") as f:
                        self.trade_history = json.load(f)
                    logger.info(f"Loaded {len(self.trade_history)} trades from history")
                except Exception as e:
                    logger.error(f"Error loading trade history: {e}")
            
            # Start balance monitoring in a separate task
            logger.info("Starting balance monitoring...")
            balance_monitor_task = asyncio.create_task(self.start_balance_monitoring())
            
            # Wait a bit to prevent too many connections at once
            await asyncio.sleep(2.0)
            
            # Start monitoring for MEV opportunities
            logger.info("Starting MEV opportunity monitoring...")
            ray_monitor_task = asyncio.create_task(self.monitor_ray_logs())
            
            # Print ready message
            logger.info("="*80)
            logger.info("✅ Bot initialized and running")
            logger.info("="*80 + "\n")
            
            # Start profit stats periodic logging
            async def log_stats_periodically():
                while True:
                    await asyncio.sleep(60 * 15)  # Log every 15 minutes
                    self.log_profit_stats()
            
            stats_task = asyncio.create_task(log_stats_periodically())
            
            # Wait for all tasks to complete (they should run indefinitely)
            await asyncio.gather(balance_monitor_task, ray_monitor_task, stats_task)
            
        except KeyboardInterrupt:
            logger.info("Trader stopped by user.")
        except Exception as e:
            logger.error(f"Fatal error: {e}")
            traceback.print_exc()
            
            if self.auto_restart and self.restart_attempts < self.max_restart_attempts:
                self.restart_attempts += 1
                logger.warning(f"Auto-restarting (attempt {self.restart_attempts}/{self.max_restart_attempts})...")
                await asyncio.sleep(ERROR_COOLDOWN)
                return await self.main()
            
        finally:
            # Log final profit stats
            self.log_profit_stats()
            logger.info("Shutting down MEV trader...")

if __name__ == "__main__":
    # Verify environment
    if not load_env_file(".env"):
        logger.warning("No .env file found, using environment variables")
    
    # Make sure the keypair exists
    keypair_path = config.get("keypair_path")
    if not keypair_path:
        logger.error("KEYPAIR_PATH not found in config or environment")
        sys.exit(1)
        
    if not os.path.exists(keypair_path):
        logger.error(f"Keypair file not found: {keypair_path}")
        sys.exit(1)
    
    # Print rate limit information
    print("\n" + "="*80)
    print("⚠️  RATE LIMITING SETTINGS")
    print("="*80)
    print(f"• Max RPC calls per second: {2}")
    print(f"• Minimum trade interval: {MIN_TRADE_INTERVAL} seconds")
    print(f"• Added artificial delays before RPC calls: 1.0 second")
    print(f"• Using exponential backoff for 429 errors")
    if isinstance(config.get("rpc_urls"), list) and len(config.get("rpc_urls")) > 1:
        print(f"• Load balancing across {len(config.get('rpc_urls'))} RPC endpoints")
    print("="*80)
    
    print("💰 PROFIT SETTINGS")
    print("="*80)
    print(f"• Target victim size: {MIN_VICTIM_SIZE}-{MAX_VICTIM_SIZE} SOL")
    print(f"• Front-run amount: {FRONT_RUN_AMOUNT} SOL")
    print(f"• Minimum profit percentage: {MIN_PROFIT_PCT*100:.1f}%")
    print(f"• Auto-restart on critical errors: {config.get_bool('auto_restart', True)}")
    print("="*80 + "\n")
    
    try:
        # Run the trader
        trader = MEVTrader()
        asyncio.run(trader.main())
    except KeyboardInterrupt:
        print("\n" + "="*80)
        print("👋 MEV Trader stopped by user")
        print("="*80 + "\n")
    except Exception as e:
        print(f"\n❌ Fatal error: {e}")
        traceback.print_exc()
        sys.exit(1) fix the spacing please write the same script just with the proper spacing please