Okay, I need to figure out why the user's MEV bot isn't finding any valid tokens and failing all transactions. Let's look at the logs they provided.

First, the bot starts up and connects to the RPC and WebSocket. The initial balance check shows they have some SOL. Then, when processing transactions, the logs repeatedly show messages like "Token X does not exist or is not a mint".

Looking at the code, the `TokenSecurity` class checks if a token exists by querying the account info. If the account doesn't exist or isn't a mint, it marks it as invalid. The logs indicate that all tokens found in transactions are being rejected for this reason.

Possible issues could be:

1. **Incorrect RPC Configuration**: If the RPC endpoint isn't responding correctly or is rate-limited, the account info checks might fail. The user is using a Helius RPC, which might have issues or require specific headers like an API key.

2. **Invalid Transaction Parsing**: The transaction logs might not be parsed correctly, leading to extracting wrong token mints. The `extract_token_mint` method might be picking up invalid addresses.

3. **Token Validation Logic Flaws**: The `verify_token_exists` method might have a bug. For example, it checks if the account data's parsed type is "mint", but maybe the RPC response format is different.

4. **Network Issues**: The Solana network might have changed, or the token mints in recent transactions are indeed invalid (e.g., new tokens that haven't been properly initialized).

Looking at the code's `verify_token_exists` function: it sends a getAccountInfo request. If the result exists and the parsed data type is "mint", it returns true. But if the RPC doesn't return that data correctly (maybe due to a different encoding or structure), this check would fail.

Also, the logs show that the user's bot is subscribing to Raydium program logs. If those logs are for failed or simulated transactions, the token mints might not be real. Alternatively, the bot might be processing testnet transactions where tokens aren't real.

Another angle: the user's RPC might not support the "jsonParsed" encoding, leading to incomplete account info. Maybe the response doesn't include the parsed data, so the check fails.

Possible fixes:

- Add more logging in `verify_token_exists` to see the raw RPC response.

- Ensure the RPC URL and API key (if required) are correctly configured.

- Check if the token mints extracted are valid by querying them manually using a Solana explorer or CLI.

- Handle different response structures from the RPC, maybe some mints are SPL tokens but not detected properly.

Also, the user's balance is low (0.21 SOL), which might affect the ability to front-run, but the logs show transactions are rejected before even attempting trades, so balance isn't the main issue here.

The key problem seems to be that all tokens are being marked as non-existent. The root cause is likely in the `verify_token_exists` method not correctly identifying valid mints, possibly due to RPC response issues or incorrect parsing.

The issue stems from invalid token mint detection due to RPC response parsing errors. Let's fix this with enhanced validation:

    Improve Token Existence Check:

python

async def verify_token_exists(self, token_mint):
    """Verify token existence with better error handling"""
    try:
        async with aiohttp.ClientSession() as session:
            response = await session.post(
                self.rpc_url,
                json={
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "getAccountInfo",
                    "params": [
                        token_mint,
                        {"encoding": "base64"}  # Use raw base64 encoding
                    ]
                },
                headers={"Content-Type": "application/json"}
            )
            data = await response.json()
            
            # Check if account exists
            if not data.get("result", {}).get("value"):
                return False
                
            # Check account is executable (mints shouldn't be)
            if data["result"]["value"]["executable"]:
                return False
                
            # Check minimum data size for SPL token
            if len(data["result"]["value"]["data"][0]) < 82:
                return False
                
            return True
            
    except Exception as e:
        logger.error(f"Token verification error: {str(e)}")
        return False

    Enhance Security Validation:

python

async def validate_token_security(self, token_mint):
    security = TokenSecurity(token_mint=token_mint)
    
    # Quick WSOL check
    if token_mint == WSOL_ADDRESS:
        security.valid = True
        return security

    # Verify account existence
    if not await self.verify_token_exists(token_mint):
        security.exists = False
        return security

    # Get mint info through token program
    try:
        mint_info = await self.trader.get_mint_info(token_mint)
        security.mint_authority_revoked = mint_info["mint_authority"] is None
    except Exception:
        return security

    # Get associated pool
    pool_address = await self.find_raydium_pool(token_mint)
    if not pool_address:
        return security

    # Check pool liquidity
    pool_data = await self.get_pool_data(pool_address)
    if pool_data:
        security.lp_locked = pool_data["lp_locked"]
        security.exit_liquidity_sol = pool_data["exit_liquidity"]

    security.valid = all([
        security.mint_authority_revoked,
        security.lp_locked,
        security.exit_liquidity_sol >= MIN_LIQUIDITY
    ])
    
    return security

    Update Raydium Pool Discovery:

python

async def find_raydium_pool(self, token_mint):
    """Find Raydium pool using more reliable method"""
    try:
        # Use Raydium API for pool discovery
        async with aiohttp.ClientSession() as session:
            async with session.get(
                f"https://api.raydium.io/v2/sdk/liquidity/mainnet.json"
            ) as response:
                pools = await response.json()
                for pool in pools["official"]:
                    if pool["base_mint"] == token_mint:
                        return pool["id"]
    except Exception:
        return await self.get_token_pool(token_mint)  # Fallback to RPC method

Critical Fixes Needed:

    Add proper error handling for RPC requests

    Update the token verification to use direct account checks instead of parsed data

    Implement fallback pool discovery methods

    Add retry logic for RPC calls

    Verify network connection and RPC endpoint status

These changes will improve token validation reliability and handle RPC inconsistencies more effectively.#!/usr/bin/env python3
"""
Quantum Raydium MEV Trader - Production Implementation
------------------------------------------------------
High-performance MEV trader for Raydium with GPU acceleration,
microsecond-level latency, and comprehensive security verification.
"""

import asyncio
import base58
import base64
import json
import logging
import os
import signal
import sys
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Any
import re
import ssl
import aiohttp
import websockets
from dotenv import load_dotenv

# Import GPU acceleration
try:
    import torch
    import numpy as np
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    logging.warning("PyTorch not installed. Running without GPU acceleration.")

# Import Rust trading core
from solana_rust_bot import SolanaTrader

# Fix for Windows event loop
if sys.platform == 'win32':
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

# Configure logging
os.makedirs("logs", exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    handlers=[
        logging.FileHandler(f"logs/solana_mev_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Constants
LAMPORTS_PER_SOL = 1_000_000_000

@dataclass
class TokenSecurity:
    mint_authority_revoked: bool = False
    lp_locked: bool = False
    has_exit_liquidity: bool = False
    exit_liquidity_sol: float = 0.0
    pool_address: str = ""
    token_mint: str = ""
    valid: bool = False
    exists: bool = False  # Added field to track if token account exists
    
@dataclass
class TradingStats:
    transactions_seen: int = 0
    raydium_swaps_found: int = 0
    viable_opportunities: int = 0
    trades_attempted: int = 0
    trades_completed: int = 0
    failed_front_runs: int = 0
    failed_back_runs: int = 0
    total_profit: float = 0.0
    total_loss: float = 0.0
    start_time: float = field(default_factory=time.time)
    token_profits: Dict[str, float] = field(default_factory=dict)
    victim_confirmed_count: int = 0
    victim_missed_count: int = 0
    initial_sol_balance: float = 0.0
    rpc_calls: int = 0
    cached_tokens_used: int = 0
    ws_connections: int = 0
    ws_messages_received: int = 0
    ws_errors: int = 0
    ws_reconnects: int = 0
    invalid_tokens_count: int = 0  # Track rejected tokens
    account_not_found_count: int = 0  # Track AccountNotFound errors
    
@dataclass
class TradeOpportunity:
    token_mint: str
    amount_sol: float
    tx_signature: str
    accounts: Dict[str, str]
    front_amount: float
    start_time: float = field(default_factory=time.time)
    front_tx: Optional[str] = None
    back_tx: Optional[str] = None
    status: str = "pending"
    profit: float = 0.0
    victim_confirmed: bool = False
    token_balance: float = 0.0
    token_decimals: int = 9
    start_sol_balance: float = 0.0
    start_wsol_balance: float = 0.0
    end_sol_balance: float = 0.0
    end_wsol_balance: float = 0.0

class RaydiumDataExtractor:
    """Extract and validate token data from Raydium pools and transaction logs"""

    def __init__(self, tensor_size=176):
        self.tensor_size = tensor_size
        
        # Initialize GPU acceleration if available
        if HAS_TORCH and torch.cuda.is_available():
            self.device = torch.device('cuda')
            # Setup CUDA for maximum performance
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False
            # Create stream for parallel processing
            self.stream = torch.cuda.Stream(priority=-3)
            self.warmup_gpu()
            logger.info(f"GPU acceleration enabled: {torch.cuda.get_device_name()}")
        else:
            self.device = None
            logger.info("Running without GPU acceleration")
            
        # Program IDs to exclude when identifying token mints
        self.program_ids = {
            os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8"),
            os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
            os.getenv("SERUM_PROGRAM_ID", "9xQeWvG816bUx9EPjHmaT23yvVM2ZWbrrpZb9PusVFin"),
            os.getenv("SYSTEM_PROGRAM_ID", "11111111111111111111111111111111"),
            os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112"),
            os.getenv("ASSOCIATED_TOKEN_PROGRAM_ID", "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"),
        }
        
    def warmup_gpu(self):
        """Warm up GPU to ensure consistent performance"""
        if not HAS_TORCH or not torch.cuda.is_available():
            return
            
        # Create tensors for warmup
        with torch.cuda.stream(self.stream):
            warmup_tensor = torch.randn(self.tensor_size, device=self.device)
            for _ in range(2500):  # Warm up cycles from your benchmarks
                result = torch.matmul(warmup_tensor, warmup_tensor)
                
        # Synchronize to ensure warmup is complete
        torch.cuda.synchronize()
        logger.info("GPU warmup completed")
        
    def decode_ray_log(self, log_str):
        """Decode Raydium ray_log with optimized parsing"""
        if "ray_log:" not in log_str:
            return None
            
        try:
            ray_log_part = log_str.split("ray_log:")[1].strip()
            
            # Multi-strategy decoding with validation
            raw_data = None
            decode_attempts = [
                lambda x: base64.b64decode(x),  # Standard
                lambda x: base64.b64decode(x + "=" * (4 - len(x) % 4) % 4),  # Add padding
                lambda x: base64.b64decode(x.replace("-", "+").replace("_", "/")),  # URL-safe
                lambda x: base64.b64decode(x.replace("-", "+").replace("_", "/") + "=" * (4 - len(x.replace("-", "+").replace("_", "/")) % 4) % 4)  # URL-safe with padding
            ]
            
            # Try all decoding strategies
            for decode_func in decode_attempts:
                try:
                    raw_data = decode_func(ray_log_part)
                    if raw_data and len(raw_data) >= 9:
                        break
                except Exception:
                    continue
                    
            if not raw_data or len(raw_data) < 9:
                return None
            
            # First byte is instruction type
            instruction_type = raw_data[0]
            
            # Next 8 bytes are the amount in little-endian
            amount_bytes = raw_data[1:9]
            amount = int.from_bytes(amount_bytes, byteorder="little") / LAMPORTS_PER_SOL
            
            # Based on instruction type, determine if buy or sell
            is_buy = instruction_type in [3, 9]  # 3 & 9 are buy instructions
            is_sell = instruction_type == 11  # 11 is sell instruction
            
            # Extract pubkeys if available
            pubkeys = []
            if len(raw_data) > 9:
                account_data = raw_data[9:]
                
                for i in range(0, len(account_data) - 31, 32):
                    try:
                        pubkey_bytes = account_data[i:i+32]
                        # Skip zero or all-one pubkeys
                        if all(b == 0 for b in pubkey_bytes) or all(b == 255 for b in pubkey_bytes):
                            continue
                            
                        pubkey = base58.b58encode(pubkey_bytes).decode('utf-8')
                        if self.is_valid_pubkey(pubkey):
                            pubkeys.append(pubkey)
                    except Exception:
                        continue
            
            return {
                "instruction_type": instruction_type,
                "amount": amount,
                "is_buy": is_buy,
                "is_sell": is_sell,
                "pubkeys": pubkeys,
                "raw_data": raw_data
            }
            
        except Exception as e:
            logger.error(f"Error decoding ray_log: {e}")
            return None
            
    def is_valid_pubkey(self, address):
        """Validate a Solana address format"""
        if not address or len(address) < 32 or len(address) > 44:
            return False
            
        try:
            # Check for concatenated addresses (e.g. 11111111... + another address)
            if address.startswith("11111111") and len(address) > 8:
                return False
                
            # Make sure we decode to 32 bytes exactly
            decoded = base58.b58decode(address)
            return len(decoded) == 32
        except Exception:
            return False
    
    def extract_token_mint(self, logs):
        """Extract token mint from logs with multi-pass approach and better validation"""
        # Method 1: Extract from Transfer logs targeting Mint Authority
        token_program_id = os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
        token_candidates = []
        
        for log in logs:
            if token_program_id in log and "Instruction: Transfer" in log:
                # Look for the next few logs for pubkeys
                token_index = logs.index(log)
                for i in range(token_index, min(token_index + 5, len(logs))):
                    address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                    matches = re.findall(address_pattern, logs[i])
                    for match in matches:
                        if match not in self.program_ids and self.is_valid_pubkey(match):
                            token_candidates.append(match)
        
        # Method 2: Extract from ray_log pubkeys
        for log in logs:
            if "ray_log:" in log:
                ray_data = self.decode_ray_log(log)
                if ray_data and ray_data["pubkeys"]:
                    for pubkey in ray_data["pubkeys"]:
                        # Skip known program IDs
                        if pubkey not in self.program_ids and self.is_valid_pubkey(pubkey):
                            token_candidates.append(pubkey)
        
        # Method 3: Extract from MintTo or related token instructions
        for log in logs:
            if token_program_id in log and "Program data:" in log and "MintTo" in log:
                address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                matches = re.findall(address_pattern, log)
                for match in matches:
                    if match not in self.program_ids and self.is_valid_pubkey(match):
                        token_candidates.append(match)
        
        # Count occurrences and pick the most likely token mint
        if token_candidates:
            # Filter out common program IDs and system accounts
            filtered_candidates = [c for c in token_candidates if not c.startswith("1111")]
            
            if filtered_candidates:
                # Count occurrences and pick the most frequent
                token_counts = {}
                for token in filtered_candidates:
                    token_counts[token] = token_counts.get(token, 0) + 1
                
                # Sort by frequency and return the most common
                sorted_tokens = sorted(token_counts.items(), key=lambda x: x[1], reverse=True)
                return sorted_tokens[0][0]
        
        return None
        
    def extract_accounts(self, logs):
        """Extract account information from logs"""
        accounts = {
            "token_program_id": os.getenv("TOKEN_PROGRAM_ID", "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
            "serum_program_id": os.getenv("SERUM_PROGRAM_ID", "9xQeWvG816bUx9EPjHmaT23yvVM2ZWbrrpZb9PusVFin"),
            "raydium_program_id": os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8"),
            "system_program_id": os.getenv("SYSTEM_PROGRAM_ID", "11111111111111111111111111111111")
        }
        
        # Account field mapping from logs
        account_fields = {
            "amm id": "amm_id",
            "authority": "amm_authority",
            "open orders": "amm_open_orders",
            "target orders": "amm_target_orders",
            "pool coin": "pool_coin_token_account",
            "pool pc": "pool_pc_token_account",
            "serum market": "serum_market",
            "serum bids": "serum_bids",
            "serum asks": "serum_asks",
            "event queue": "serum_event_queue",
            "coin vault": "serum_coin_vault",
            "pc vault": "serum_pc_vault",
            "vault signer": "serum_vault_signer"
        }
        
        # Extract accounts from logs
        for log in logs:
            if ":" not in log:
                continue
                
            # Check for account patterns
            for pattern, field_name in account_fields.items():
                if pattern in log.lower():
                    parts = log.split(":")
                    if len(parts) < 2:
                        continue
                        
                    # Extract pubkey after pattern
                    address_part = parts[1].strip()
                    address_pattern = r'\b[1-9A-HJ-NP-Za-km-z]{32,44}\b'
                    matches = re.findall(address_pattern, address_part)
                    
                    for match in matches:
                        if self.is_valid_pubkey(match):
                            accounts[field_name] = match
                            break
        
        # Extract AMM ID from ray_log as fallback
        if "amm_id" not in accounts:
            for log in logs:
                if "ray_log:" in log:
                    ray_data = self.decode_ray_log(log)
                    if ray_data and ray_data["pubkeys"] and len(ray_data["pubkeys"]) >= 5:
                        accounts["amm_id"] = ray_data["pubkeys"][0]
                        break
        
        return accounts

class GPUAcceleratedTokenValidator:
    """Token validator with GPU acceleration and enhanced security checks"""
    
    def __init__(self, trader, rpc_url):
        self.trader = trader
        self.rpc_url = rpc_url
        self.security_cache = {}
        
        # Security profile for checking binary layout
        self.security_profiles = [
            {"mint_authority": 72, "lp_lock": 88, "liquidity": (128, 136)},
            {"mint_authority": 70, "lp_lock": 80, "liquidity": (128, 136)},
            {"mint_authority": 64, "lp_lock": 100, "liquidity": (128, 136)},
        ]
        
    async def verify_token_exists(self, token_mint):
        """Verify a token mint address exists before validation"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.rpc_url,
                    json={
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "getAccountInfo",
                        "params": [
                            token_mint,
                            {"encoding": "jsonParsed"}
                        ]
                    },
                    headers={"Content-Type": "application/json"}
                ) as response:
                    data = await response.json()
                    
                    if "result" in data and data["result"] and data["result"]["value"]:
                        # Check if it's actually a token mint by looking for mint info
                        account_data = data["result"]["value"]
                        if "data" in account_data and "parsed" in account_data["data"]:
                            parsed = account_data["data"]["parsed"]
                            if "type" in parsed and parsed["type"] == "mint":
                                return True
            
            return False
        except Exception as e:
            logger.error(f"Error verifying token {token_mint[:8]} exists: {e}")
            return False
            
    async def validate_token_security(self, token_mint):
        """Validate token security features comprehensively with existence check"""
        # Create security result
        security = TokenSecurity(token_mint=token_mint)
        
        # Skip WSOL - always safe
        if token_mint == os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112"):
            return TokenSecurity(
                mint_authority_revoked=True,
                lp_locked=True,
                has_exit_liquidity=True,
                exit_liquidity_sol=1000000.0,
                pool_address="NATIVE_SOL",
                token_mint=token_mint,
                valid=True,
                exists=True
            )
        
        # Check cache
        cache_time = int(os.getenv("SECURITY_CACHE_SECONDS", "600"))
        if token_mint in self.security_cache:
            last_check, security = self.security_cache[token_mint]
            if time.time() - last_check < cache_time:
                return security
        
        try:
            # 1. First verify the token actually exists and is a mint
            token_exists = await self.verify_token_exists(token_mint)
            if not token_exists:
                logger.info(f"Token {token_mint[:8]} validation failed: token does not exist or is not a mint")
                security.exists = False
                self.security_cache[token_mint] = (time.time(), security)
                return security
                
            security.exists = True
            
            # 2. Check mint authority
            security.mint_authority_revoked = self.trader.verify_token_safety(token_mint)
            if not security.mint_authority_revoked:
                logger.info(f"Token {token_mint[:8]} validation failed: mint authority not revoked")
                self.security_cache[token_mint] = (time.time(), security)
                return security
            
            # 3. Get pool address and check liquidity with RPC
            pool_address = await self.get_token_pool(token_mint)
            if not pool_address:
                logger.info(f"Token {token_mint[:8]} validation failed: no pool found")
                self.security_cache[token_mint] = (time.time(), security)
                return security
                
            security.pool_address = pool_address
            
            # 4. Get pool data and validate security
            pool_data = await self.get_pool_data(pool_address)
            if not pool_data:
                logger.info(f"Token {token_mint[:8]} validation failed: failed to get pool data")
                self.security_cache[token_mint] = (time.time(), security)
                return security
                
            # 5. Analyze pool binary data
            security_result = self.analyze_pool_security(pool_data)
            security.lp_locked = security_result.get("lp_locked", False)
            security.exit_liquidity_sol = security_result.get("exit_liquidity", 0.0)
            security.has_exit_liquidity = security.exit_liquidity_sol >= float(os.getenv("MIN_EXIT_LIQUIDITY_SOL", "10.0"))
            
            # 6. Final validity check
            security.valid = (
                security.mint_authority_revoked and 
                security.lp_locked and 
                security.has_exit_liquidity
            )
            
            # Cache result
            self.security_cache[token_mint] = (time.time(), security)
            
            # Log result
            logger.info(
                f"Token {token_mint[:8]} validation: " +
                f"mint_revoked={security.mint_authority_revoked}, " +
                f"lp_locked={security.lp_locked}, " +
                f"exit_liquidity={security.exit_liquidity_sol:.2f} SOL, " +
                f"valid={security.valid}"
            )
            
            return security
        except Exception as e:
            logger.error(f"Error validating token {token_mint[:8]}: {e}")
            security.exists = False
            self.security_cache[token_mint] = (time.time(), security)
            return security
    
    async def get_token_pool(self, token_mint):
        """Get the pool address for a token"""
        try:
            # Use RPC to find the pool containing this token
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.rpc_url,
                    json={
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "getProgramAccounts",
                        "params": [
                            os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8"),
                            {
                                "encoding": "base64",
                                "filters": [
                                    {"dataSize": 680},  # Raydium V4 pool size
                                    {"memcmp": {"offset": 72, "bytes": token_mint}}
                                ]
                            }
                        ]
                    },
                    headers={"Content-Type": "application/json"}
                ) as response:
                    data = await response.json()
                    
                    if "result" in data and len(data["result"]) > 0:
                        return data["result"][0]["pubkey"]
            
            return None
        except Exception as e:
            logger.error(f"Error finding pool for token {token_mint[:8]}: {e}")
            return None
    
    async def get_pool_data(self, pool_address):
        """Get the raw binary data for a pool"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.rpc_url,
                    json={
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "getAccountInfo",
                        "params": [
                            pool_address,
                            {"encoding": "base64"}
                        ]
                    },
                    headers={"Content-Type": "application/json"}
                ) as response:
                    data = await response.json()
                    
                    if "result" in data and data["result"] and "value" in data["result"]:
                        value = data["result"]["value"]
                        if "data" in value and len(value["data"]) > 0:
                            return base64.b64decode(value["data"][0])
            
            return None
        except Exception as e:
            logger.error(f"Error getting pool data for {pool_address}: {e}")
            return None
    
    def analyze_pool_security(self, data):
        """Analyze pool security features"""
        results = {
            "exit_liquidity": 0.0,
            "lp_locked": False
        }
        
        # Try all security profiles
        for profile in self.security_profiles:
            try:
                # Check liquidity
                if len(data) >= profile["liquidity"][1]:
                    liquidity_bytes = data[profile["liquidity"][0]:profile["liquidity"][1]]
                    liquidity_raw = int.from_bytes(liquidity_bytes, "little")
                    # Convert to SOL amount using scaling factor
                    results["exit_liquidity"] = liquidity_raw / 1e6
                
                # Check LP lock
                if len(data) >= profile["lp_lock"] + 8:
                    lp_raw_value = int.from_bytes(data[profile["lp_lock"]:profile["lp_lock"]+8], "little")
                    results["lp_locked"] = lp_raw_value > 0
                
                # If we found valid data, exit loop
                if results["exit_liquidity"] > 0:
                    break
                    
            except Exception:
                continue
                
        return results

class RaydiumMEVTrader:
    """High-performance MEV trader for Raydium with GPU acceleration"""
    
    def __init__(self):
        # Load environment variables
        load_dotenv()
        
        # Create thread pool for parallel processing
        self.thread_pool = ThreadPoolExecutor(
            max_workers=int(os.getenv("THREADS", "48"))
        )
        
        # Initialize trader with keypair
        self.trader = SolanaTrader(
            rpc_url=os.getenv("RPC_URL", "https://winny-rychu7-fast-mainnet.helius-rpc.com"),
            ws_url=os.getenv("WS_URL", "wss://winny-rychu7-fast-mainnet.helius-rpc.com"),
            keypair_path=os.getenv("KEYPAIR_PATH", "C:\\solana_rust_bot\\keypair.bin"),
            compute_limit=int(os.getenv("COMPUTE_LIMIT", "50000"))
        )
        
        # Initialize wallet info
        self.wallet_address = self.trader.get_address()
        
        # Initialize components
        self.token_extractor = RaydiumDataExtractor(
            tensor_size=int(os.getenv("TENSOR_SIZE", "176"))
        )
        self.token_validator = GPUAcceleratedTokenValidator(
            self.trader, 
            os.getenv("RPC_URL", "https://winny-rychu7-fast-mainnet.helius-rpc.com")
        )
        
        # State tracking
        self.running = False
        self.active_opportunity = None
        self.processed_txs = set()
        self.valid_tokens = set()
        self.invalid_tokens = set()
        self.last_balance_update = 0
        self.sol_balance = 0
        self.wsol_balance = 0
        self.wsol_account = os.getenv("WSOL_TOKEN_ACCOUNT", "")
        self.stats = TradingStats()
        self.balance_lock = asyncio.Lock()
        
        # Config
        self.front_run_amount = float(os.getenv("FRONT_RUN_AMOUNT", "0.05"))  # Reduced from 0.3 to 0.05
        self.min_victim_size = float(os.getenv("SANDWICH_MIN_VICTIM_SIZE", "3.0"))
        self.max_victim_size = float(os.getenv("SANDWICH_MAX_VICTIM_SIZE", "320.0"))
        self.min_sol_balance = float(os.getenv("MIN_SOL_BALANCE", "0.01"))  # Reduced from 0.03 to 0.01
        self.min_wsol_balance = float(os.getenv("MIN_WSOL_BALANCE", "0.002"))
        self.slippage_bps = int(os.getenv("SLIPPAGE_BPS", "1000"))
        self.min_profit = float(os.getenv("MIN_PROFIT", "0.002"))
        self.wsol_wrap_interval = int(os.getenv("WSOL_WRAP_INTERVAL", "30"))
        self.last_wrap_time = 0
        
        logger.info(f"Initialized MEV trader with wallet: {self.wallet_address}")
        
    async def run(self):
        """Run the MEV trader"""
        self.running = True
        
        # Setup signal handlers
        for sig in (signal.SIGINT, signal.SIGTERM):
            signal.signal(sig, self._handle_exit)
            
        logger.info("Starting Raydium MEV Trader...")
        
        # Initial balance check and WSOL account creation
        try:
            # Verify WSOL account exists or create it
            await self.ensure_wsol_account_exists()
            
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            self.last_balance_update = time.time()
            
            self.stats.initial_sol_balance = self.sol_balance + self.wsol_balance
            logger.info(f"Initial wallet balance: {self.stats.initial_sol_balance:.4f} SOL")
        except Exception as e:
            logger.error(f"Error getting initial balance: {e}")
        
        # Start tasks
        tasks = [
            asyncio.create_task(self.monitor_mempool()),
            asyncio.create_task(self.monitor_balances())
        ]
        
        # Keep the bot running
        while self.running:
            await asyncio.sleep(0.1)
            await self.check_active_opportunity()
        
        # Safe shutdown
        await self.safe_shutdown()
        self.print_stats()
        
        # Cancel tasks
        for task in tasks:
            task.cancel()
    
    async def ensure_wsol_account_exists(self):
        """Ensure WSOL token account exists, creating it if needed"""
        wsol_address = os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
        
        # Check if WSOL account is provided in env
        if not self.wsol_account:
            try:
                # Get the default WSOL token account from trader
                self.wsol_account = self.trader.get_associated_token_address(
                    self.wallet_address, wsol_address
                )
                logger.info(f"Using calculated WSOL account: {self.wsol_account}")
            except Exception as e:
                logger.error(f"Error getting WSOL account: {e}")
                return
        
        # Check if the account already exists
        wsol_balance = -1
        try:
            wsol_balance = self.trader.get_token_balance(wsol_address)
        except Exception:
            logger.info("WSOL account does not exist, creating it...")
            
        # If balance check failed, create the WSOL account
        if wsol_balance < 0:
            try:
                # Create the token account
                result = self.trader.create_token_account_if_needed(wsol_address)
                if result:
                    logger.info(f"Successfully created WSOL account: {self.wsol_account}")
                    # Wrap a small amount to initialize the account
                    if self.sol_balance > 0.01:
                        logger.info("Wrapping 0.01 SOL to initialize WSOL account")
                        wrap_result = self.trader.wrap_sol(0.01)
                        if wrap_result:
                            logger.info("Successfully initialized WSOL account")
                            self.last_wrap_time = time.time()
            except Exception as e:
                logger.error(f"Error creating WSOL account: {e}")
    
    def _handle_exit(self, sig, frame):
        """Handle exit signals gracefully"""
        if not self.running:
            return
            
        logger.info("Shutdown signal received, cleaning up...")
        self.running = False
        
    async def safe_shutdown(self):
        """Clean shutdown with completion of active trades"""
        logger.info("Initiating safe shutdown...")
        
        # Complete active opportunity if any
        if self.active_opportunity and self.active_opportunity.status == "frontrun":
            logger.info(f"Completing active opportunity before shutdown")
            await self.execute_emergency_sell(self.active_opportunity)
        
        # Unwrap excess WSOL
        if self.wsol_balance > self.min_wsol_balance * 2:
            await self.unwrap_wsol()
            
        logger.info("Shutdown complete")
        
    def print_stats(self):
        """Print detailed trading statistics"""
        duration = time.time() - self.stats.start_time
        hours = duration // 3600
        minutes = (duration % 3600) // 60
        seconds = duration % 60
        
        logger.info("=" * 50)
        logger.info("MEV Trading Statistics:")
        logger.info(f"Runtime: {int(hours)}h {int(minutes)}m {int(seconds)}s")
        logger.info(f"Transactions seen: {self.stats.transactions_seen}")
        logger.info(f"Raydium swaps found: {self.stats.raydium_swaps_found}")
        logger.info(f"Viable opportunities: {self.stats.viable_opportunities}")
        logger.info(f"Trades attempted: {self.stats.trades_attempted}")
        logger.info(f"Trades completed: {self.stats.trades_completed}")
        logger.info(f"Failed front-runs: {self.stats.failed_front_runs}")
        logger.info(f"Failed back-runs: {self.stats.failed_back_runs}")
        logger.info(f"Total profit: {self.stats.total_profit:.6f} SOL")
        logger.info(f"Total loss: {self.stats.total_loss:.6f} SOL")
        logger.info(f"Net profit: {self.stats.total_profit - self.stats.total_loss:.6f} SOL")
        
        # Token validation stats
        logger.info(f"Invalid tokens rejected: {self.stats.invalid_tokens_count}")
        logger.info(f"AccountNotFound errors: {self.stats.account_not_found_count}")
        
        # Calculate success rate
        if self.stats.trades_attempted > 0:
            success_rate = (self.stats.trades_completed / self.stats.trades_attempted) * 100
            logger.info(f"Success rate: {success_rate:.2f}%")
            
        # Calculate profit per completed trade
        if self.stats.trades_completed > 0:
            profit_per_trade = (self.stats.total_profit - self.stats.total_loss) / self.stats.trades_completed
            logger.info(f"Average profit per completed trade: {profit_per_trade:.6f} SOL")
        
        # WebSocket stats
        logger.info(f"WebSocket connections: {self.stats.ws_connections}")
        logger.info(f"WebSocket messages received: {self.stats.ws_messages_received}")
        logger.info(f"WebSocket errors: {self.stats.ws_errors}")
        logger.info(f"WebSocket reconnects: {self.stats.ws_reconnects}")
        
        logger.info("=" * 50)
    
    async def monitor_mempool(self):
        """Monitor Raydium program logs for MEV opportunities"""
        logger.info(f"Starting mempool monitoring via WebSocket")
        
        retry_count = 0
        backoff = 1.0
        
        while self.running:
            try:
                # Setup WebSocket connection
                headers = {"Content-Type": "application/json"}
                api_key = os.getenv("API_KEY", "")
                if api_key:
                    headers["api-key"] = api_key
                
                # Create a custom SSL context that doesn't verify
                ssl_context = ssl.create_default_context()
                ssl_context.check_hostname = False
                ssl_context.verify_mode = ssl.CERT_NONE
                
                ws_url = os.getenv("WS_URL", "wss://winny-rychu7-fast-mainnet.helius-rpc.com")
                
                # Connect to WebSocket
                async with websockets.connect(
                    ws_url,
                    extra_headers=headers,
                    ping_interval=30,
                    ping_timeout=20,
                    close_timeout=10,
                    max_size=10 * 1024 * 1024,
                    ssl=ssl_context
                ) as ws:
                    self.stats.ws_connections += 1
                    logger.info(f"Connected to WebSocket: {ws_url}")
                    
                    # Reset retry count on successful connection
                    retry_count = 0
                    backoff = 1.0
                    
                    # Subscribe to Raydium program logs
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0", 
                        "id": 1,
                        "method": "logsSubscribe",
                        "params": [
                            {"mentions": [os.getenv("RAYDIUM_PROGRAM_ID", "675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8")]},
                            {"commitment": "confirmed"}
                        ]
                    }))
                    
                    # Wait for logs subscription confirmation
                    logs_response = await ws.recv()
                    self.stats.ws_messages_received += 1
                    logs_data = json.loads(logs_response)
                    
                    if "result" not in logs_data:
                        error_msg = logs_data.get("error", {}).get("message", "Unknown error")
                        logger.error(f"Failed to subscribe to logs: {error_msg}")
                        await asyncio.sleep(backoff)
                        backoff = min(30, backoff * 1.5)
                        continue
                    
                    logs_subscription_id = logs_data["result"]
                    logger.info(f"Successfully subscribed to Raydium logs (ID: {logs_subscription_id})")
                    
                    # Subscribe to account updates for balance tracking
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0", 
                        "id": 2,
                        "method": "accountSubscribe",
                        "params": [
                            self.wallet_address,
                            {"commitment": "confirmed", "encoding": "jsonParsed"}
                        ]
                    }))
                    
                    # Wait for account subscription confirmation
                    account_response = await ws.recv()
                    self.stats.ws_messages_received += 1
                    account_data = json.loads(account_response)
                    
                    if "result" in account_data:
                        account_subscription_id = account_data["result"]
                        logger.info(f"Successfully subscribed to wallet account (ID: {account_subscription_id})")
                    else:
                        logger.warning(f"Failed to subscribe to wallet account: {account_data}")
                    
                    # Set up timing variables
                    last_ping_time = time.time()
                    last_heartbeat = time.time()
                    
                    # Process WebSocket messages
                    while self.running:
                        try:
                            # Set shorter timeout to allow for periodic checks
                            message = await asyncio.wait_for(ws.recv(), timeout=5)
                            self.stats.ws_messages_received += 1
                            
                            # Reset heartbeat timer on each message
                            last_heartbeat = time.time()
                            
                            data = json.loads(message)
                            
                            # Handle log notifications (Raydium transactions)
                            if ("method" in data and data["method"] == "logsNotification" and 
                                  "params" in data and "subscription" in data["params"] and 
                                  data["params"]["subscription"] == logs_subscription_id):
                                await self.process_logs_notification(data)
                                
                            # Handle account notifications (balance updates)
                            elif ("method" in data and data["method"] == "accountNotification" and 
                                  "params" in data and "subscription" in data["params"]):
                                await self.update_wallet_balance()
                            
                        except asyncio.TimeoutError:
                            # Check if heartbeat needed
                            current_time = time.time()
                            time_since_last_message = current_time - last_heartbeat
                            
                            # If no messages for twice the heartbeat interval, send ping
                            if time_since_last_message > 20:
                                logger.info(f"No messages for {time_since_last_message:.1f}s, sending heartbeat ping")
                                try:
                                    pong = await ws.ping()
                                    await pong  # Wait for pong response
                                    last_heartbeat = current_time
                                except Exception as e:
                                    logger.warning(f"Heartbeat ping failed: {e}")
                                    # Connection likely dead, break to reconnect
                                    break
                            
                            # Regular checks on timeout
                            await self.check_active_opportunity()
                            
                            # Send regular ping if needed
                            if current_time - last_ping_time > 30:
                                try:
                                    await ws.ping()
                                    last_ping_time = current_time
                                except Exception as e:
                                    logger.warning(f"Regular ping failed: {e}")
                                
                        except websockets.exceptions.ConnectionClosed as e:
                            self.stats.ws_errors += 1
                            logger.warning(f"WebSocket connection closed ({e.code}): {e.reason}, reconnecting...")
                            break
                            
                        except json.JSONDecodeError:
                            self.stats.ws_errors += 1
                            logger.warning("Received invalid JSON, ignoring message")
                            continue
                            
                        except Exception as e:
                            self.stats.ws_errors += 1
                            logger.error(f"Error processing message: {e}")
                            logger.error(traceback.format_exc())
                
                # Connection closed, now reconnect
                logger.info("WebSocket connection ended, reconnecting...")
                self.stats.ws_reconnects += 1
            
            except asyncio.TimeoutError:
                self.stats.ws_errors += 1
                logger.warning(f"WebSocket connection timeout, reconnecting...")
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            except websockets.exceptions.WebSocketException as e:
                self.stats.ws_errors += 1
                retry_count += 1
                logger.error(f"WebSocket error (retry {retry_count}): {e}")
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            except Exception as e:
                self.stats.ws_errors += 1
                retry_count += 1
                logger.error(f"Unexpected error: {e}")
                logger.error(traceback.format_exc())
                await asyncio.sleep(backoff)
                backoff = min(30, backoff * 1.5)
                
            # Always try to reconnect regardless of error
            if self.running:
                logger.info("Attempting to reconnect WebSocket...")
    
    async def update_wallet_balance(self):
        """Update wallet balance after transaction"""
        async with self.balance_lock:
            self.sol_balance = self.trader.get_sol_balance()
            self.wsol_balance = self.trader.get_token_balance(
                os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
            )
            self.last_balance_update = time.time()
            
            logger.info(f"Updated wallet balance: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
            
            # Check for victim confirmation if we have active opportunity
            if (self.active_opportunity and 
                self.active_opportunity.status == "frontrun" and 
                not self.active_opportunity.victim_confirmed):
                
                # Mark victim as confirmed and execute back-run
                self.active_opportunity.victim_confirmed = True
                self.stats.victim_confirmed_count += 1
                
                logger.info(f"Victim transaction confirmed for {self.active_opportunity.token_mint[:8]}")
                await asyncio.sleep(0.1)  # Small delay
                await self.execute_back_run(self.active_opportunity)
    
    async def process_logs_notification(self, data):
        """Process Raydium program logs for MEV opportunities"""
        try:
            # Skip if shutting down or already have active opportunity
            if not self.running or self.active_opportunity:
                return
                
            # Extract logs and signature
            params = data.get("params", {})
            result = params.get("result", {})
            value = result.get("value", {})
            
            logs = value.get("logs", [])
            tx_signature = value.get("signature", "unknown")
            
            # Skip if no logs or already processed
            if not logs or tx_signature in self.processed_txs:
                return
                
            self.stats.transactions_seen += 1
            self.processed_txs.add(tx_signature)
            
            # Limit processed_txs size
            if len(self.processed_txs) > 5000:
                self.processed_txs = set(list(self.processed_txs)[-2500:])
            
            # Check for ray_log entries
            ray_log_found = False
            token_mint = None
            accounts = {}
            amount_sol = 0.0
            is_buy = False
            
            for log in logs:
                if "ray_log:" in log:
                    ray_log_found = True
                    ray_data = self.token_extractor.decode_ray_log(log)
                    
                    if ray_data and ray_data["instruction_type"] == 3:  # We want 0x03 instructions
                        self.stats.raydium_swaps_found += 1
                        
                        # Check amount to see if it's a viable opportunity
                        amount_sol = ray_data["amount"]
                        is_buy = ray_data["is_buy"]
                        
                        if (is_buy and 
                            self.min_victim_size <= amount_sol <= self.max_victim_size):
                            
                            # Extract token mint (parallelized)
                            token_mint = await asyncio.get_event_loop().run_in_executor(
                                self.thread_pool,
                                lambda: self.token_extractor.extract_token_mint(logs)
                            )
                            
                            if token_mint:
                                # Extract accounts for swap
                                accounts = self.token_extractor.extract_accounts(logs)
                                accounts["token_mint"] = token_mint
                                
                                logger.info(f"Found potential MEV: {tx_signature[:8]}, "
                                          f"Buy {amount_sol:.4f} SOL, "
                                          f"Token: {token_mint[:8]}")
                                          
                                self.stats.viable_opportunities += 1
                                          
                                # Process opportunity
                                await self.process_opportunity(token_mint, amount_sol, tx_signature, accounts)
                        
                        else:
                            # Log if it's a Raydium swap but not a viable opportunity
                            token_info = "Unknown"
                            if ray_data["pubkeys"] and len(ray_data["pubkeys"]) > 3:
                                possible_token = ray_data["pubkeys"][3]
                                if self.token_extractor.is_valid_pubkey(possible_token):
                                    token_info = possible_token[:8]
                                    
                            logger.debug(f"Raydium swap (not viable): {tx_signature[:8]}, "
                                    f"Buy: {is_buy}, Amount: {amount_sol:.4f} SOL, "
                                    f"Token: {token_info}")
            
            if not ray_log_found:
                logger.debug(f"Transaction without ray_log: {tx_signature[:8]}")
                
        except Exception as e:
            logger.error(f"Error processing logs: {e}")
            logger.error(traceback.format_exc())
    
    async def check_active_opportunity(self):
        """Check and update active opportunity status"""
        if not self.active_opportunity:
            return
            
        # Handle timeout for victim confirmation
        timeout = int(os.getenv("VICTIM_CONFIRM_TIMEOUT", "15"))
        if (self.active_opportunity.status == "frontrun" and 
            not self.active_opportunity.victim_confirmed and
            time.time() - self.active_opportunity.start_time > timeout):
            
            logger.warning(f"Victim tx confirmation timeout for {self.active_opportunity.token_mint[:8]}")
            self.stats.victim_missed_count += 1
            
            # Execute emergency sell
            await self.execute_emergency_sell(self.active_opportunity)
    
    async def process_opportunity(self, token_mint, amount_sol, tx_signature, accounts):
        """Process an MEV opportunity"""
        try:
            # Skip if already have active opportunity or shutting down
            if self.active_opportunity or not self.running:
                return
                
            # Quick validation using cached results
            if token_mint in self.invalid_tokens:
                logger.debug(f"Skipping known invalid token: {token_mint[:8]}")
                return
                
            # Check valid tokens cache
            if token_mint in self.valid_tokens:
                is_valid = True
                logger.info(f"Using cached valid token: {token_mint[:8]}")
                self.stats.cached_tokens_used += 1
            else:
                # Verify token exists first (cheap check)
                token_exists = await self.token_validator.verify_token_exists(token_mint)
                if not token_exists:
                    logger.info(f"Token {token_mint[:8]} does not exist or is not a mint")
                    self.invalid_tokens.add(token_mint)
                    self.stats.invalid_tokens_count += 1
                    self.stats.account_not_found_count += 1
                    return
                
                # Validate token security features
                security = await self.token_validator.validate_token_security(token_mint)
                is_valid = security.valid
                
                # Update cache based on result
                if is_valid:
                    self.valid_tokens.add(token_mint)
                    logger.info(f"Validated new token: {token_mint[:8]}")
                else:
                    self.invalid_tokens.add(token_mint)
                    self.stats.invalid_tokens_count += 1
                    logger.info(f"Invalid token rejected: {token_mint[:8]}: " +
                              f"mint_revoked={security.mint_authority_revoked}, " +
                              f"lp_locked={security.lp_locked}, " +
                              f"exit_liquidity={security.exit_liquidity_sol:.2f} SOL")
                    return
            
            # Calculate front-run amount based on victim size
            # Dynamic calculation based on victim size with square root scaling
            front_amount = min(self.front_run_amount, self.sol_balance * 0.2)  # Never use more than 20% of balance
            if amount_sol > 10.0:
                # Scale up for larger victims, limited to 2x base amount
                front_amount = min(
                    self.front_run_amount * (amount_sol / 10.0) ** 0.5,
                    self.front_run_amount * 2,
                    self.sol_balance * 0.2  # Still capped at 20% of balance
                )
            
            # Create opportunity object
            opportunity = TradeOpportunity(
                token_mint=token_mint,
                amount_sol=amount_sol,
                tx_signature=tx_signature,
                accounts=accounts,
                front_amount=front_amount,
                start_time=time.time(),
                start_sol_balance=self.sol_balance,
                start_wsol_balance=self.wsol_balance
            )
            
            # Check if we have enough balance with reduced requirements
            async with self.balance_lock:
                total_balance = self.sol_balance + self.wsol_balance
                min_required = front_amount + self.min_sol_balance
                
                if total_balance < min_required:
                    logger.warning(f"Insufficient balance: {total_balance:.4f} SOL, need {min_required:.4f} SOL")
                    return
                    
                # Make sure we have enough WSOL - with reduced buffer
                wsol_ok = await self.ensure_wsol(front_amount * 1.2)  # Reduced from 1.5 to 1.2
                
                if not wsol_ok:
                    logger.error(f"Failed to ensure WSOL for {token_mint[:8]}")
                    return
                    
                # Update WSOL balance after wrapping
                self.wsol_balance = self.trader.get_token_balance(
                    os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                )
            
            # Check profitability estimation
            profit_estimate = self.estimate_profit(front_amount, amount_sol)
            if not profit_estimate["is_profitable"]:
                logger.info(f"Skipping unprofitable opportunity: {token_mint[:8]}, estimated profit: {profit_estimate['net_profit']:.6f} SOL")
                return
                
            # Set as active opportunity
            self.active_opportunity = opportunity
            self.stats.trades_attempted += 1
            
            logger.info(f"Executing front-run for victim tx {tx_signature[:8]} with {front_amount:.4f} SOL -> {token_mint[:8]}")
            
            # Execute front-run
            front_tx_signature = self.execute_swap(
                token_mint=token_mint,
                amount_sol=front_amount,
                accounts=accounts,
                is_buy=True
            )
            
            if front_tx_signature:
                logger.info(f"Front-run successful: {front_tx_signature[:8]}")
                opportunity.front_tx = front_tx_signature
                opportunity.status = "frontrun"
                
                # Get token balance and decimals for accurate back-run
                opportunity.token_balance = self.trader.get_token_balance(token_mint)
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
                
                logger.info(f"Received {opportunity.token_balance} tokens (decimals: {opportunity.token_decimals})")
                
                # Now wait for victim confirmation via WebSocket balance updates
            else:
                logger.error(f"Front-run failed: null tx")
                self.stats.failed_front_runs += 1
                self.active_opportunity = None
                
        except Exception as e:
            logger.error(f"Error processing opportunity: {e}")
            logger.error(traceback.format_exc())
            self.active_opportunity = None
    
    def execute_swap(self, token_mint, amount_sol, accounts, is_buy):
        """Execute a swap transaction with correct instruction format"""
        try:
            # Create instruction data with proper 0x09 format
            instr_data = bytearray(17)  # 17 bytes: 1 byte discriminator + 8 bytes amount + 8 bytes min_out
            
            # Set discriminator to 0x09 (Raydium V4 swap)
            instr_data[0] = 0x09
            
            # Calculate amount in lamports
            amount_lamports = int(amount_sol * LAMPORTS_PER_SOL)
            
            # Convert amount to little-endian bytes
            amount_bytes = amount_lamports.to_bytes(8, byteorder="little")
            
            # Calculate min_out with slippage
            slippage = self.slippage_bps / 10000.0  # Convert from basis points to percentage
            min_out_lamports = int(amount_lamports * (1.0 - slippage))
            min_out_bytes = min_out_lamports.to_bytes(8, byteorder="little")
            
            # Set amount_in (bytes 1-8)
            for i in range(8):
                instr_data[i+1] = amount_bytes[i]
                
            # Set min_out (bytes 9-16)
            for i in range(8):
                instr_data[i+9] = min_out_bytes[i]
            
            # Execute swap with Rust trading core
            tx_signature = self.trader.swap_with_accounts(
                instr_data,  # Proper 0x09 instruction data
                accounts.get("amm_id", ""),
                accounts.get("amm_authority", ""),
                accounts.get("amm_open_orders", ""),
                accounts.get("amm_target_orders", ""),
                accounts.get("pool_coin_token_account", ""),
                accounts.get("pool_pc_token_account", ""),
                accounts.get("serum_market", ""),
                accounts.get("serum_bids", ""),
                accounts.get("serum_asks", ""),
                accounts.get("serum_event_queue", ""),
                accounts.get("serum_coin_vault", ""),
                accounts.get("serum_pc_vault", ""),
                accounts.get("serum_vault_signer", ""),
                token_mint,
                is_buy
            )
            
            return tx_signature
        
        except Exception as e:
            logger.error(f"Error executing swap: {e}")
            return None
    
    async def execute_back_run(self, opportunity):
        """Execute back-run after victim tx confirmation"""
        try:
            token_mint = opportunity.token_mint
            
            # Get token balance with verification
            if opportunity.token_balance <= 0:
                opportunity.token_balance = self.trader.get_token_balance(token_mint)
            
            if opportunity.token_balance <= 0:
                logger.warning(f"No tokens to sell for back-run: {token_mint[:8]}")
                self.active_opportunity = None
                return
                
            logger.info(f"Back-running with {opportunity.token_balance} tokens of {token_mint[:8]}")
            
            # Get token decimals if not already known
            if opportunity.token_decimals == 9:  # Default value
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
            
            # Add optimal timing delay based on victim size
            victim_size = opportunity.amount_sol
            optimal_delay = min(0.2 + (victim_size / 100.0), 0.8)
            await asyncio.sleep(optimal_delay)
            
            # Execute back-run (sell 100% of tokens)
            back_tx_signature = self.execute_token_sell(
                token_mint=token_mint,
                token_amount=opportunity.token_balance,
                token_decimals=opportunity.token_decimals,
                accounts=opportunity.accounts
            )
            
            if back_tx_signature:
                logger.info(f"Back-run successful: {back_tx_signature[:8]}")
                opportunity.back_tx = back_tx_signature
                
                # Get updated balances
                async with self.balance_lock:
                    self.sol_balance = self.trader.get_sol_balance()
                    self.wsol_balance = self.trader.get_token_balance(
                        os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                    )
                    self.last_balance_update = time.time()
                    
                    # Calculate profit/loss
                    total_end = self.sol_balance + self.wsol_balance
                    total_start = opportunity.start_sol_balance + opportunity.start_wsol_balance
                    
                    # Calculate transaction costs
                    tx_cost = 0.000005 * 2  # Base fee for two transactions
                    compute_cost = (int(os.getenv("COMPUTE_LIMIT", "50000")) / 1000000) * (int(os.getenv("PRIORITY_FEE", "950000")) / LAMPORTS_PER_SOL) * 2
                    total_cost = tx_cost + compute_cost
                    
                    raw_profit = total_end - total_start
                    net_profit = raw_profit - total_cost
                    
                    opportunity.profit = net_profit
                    opportunity.end_sol_balance = self.sol_balance
                    opportunity.end_wsol_balance = self.wsol_balance
                    opportunity.status = "complete"
                    
                    self.stats.trades_completed += 1
                    
                    if net_profit >= 0:
                        self.stats.total_profit += net_profit
                        roi_percent = (net_profit / opportunity.front_amount) * 100
                        logger.info(f"Trade profit: +{net_profit:.6f} SOL (ROI: {roi_percent:.2f}%)")
                    else:
                        self.stats.total_loss += abs(net_profit)
                        logger.info(f"Trade loss: {net_profit:.6f} SOL")
            else:
                logger.error(f"Back-run failed: null tx")
                self.stats.failed_back_runs += 1
                await self.execute_emergency_sell(opportunity)
                return
            
            # Clear active opportunity
            self.active_opportunity = None
                
        except Exception as e:
            logger.error(f"Error in back-run: {e}")
            logger.error(traceback.format_exc())
            await self.execute_emergency_sell(opportunity)
    
    def execute_token_sell(self, token_mint, token_amount, token_decimals, accounts):
        """Sell tokens with proper instruction format"""
        try:
            # Create instruction data with proper 0x09 format
            instr_data = bytearray(17)  # 17 bytes: 1 byte discriminator + 8 bytes amount + 8 bytes min_out
            
            # Set discriminator to 0x09 (Raydium V4 swap)
            instr_data[0] = 0x09
            
            # Calculate token amount in smallest units
            token_units = int(token_amount * (10 ** token_decimals))
            
            # Convert amount to little-endian bytes
            amount_bytes = token_units.to_bytes(8, byteorder="little")
            
            # Calculate min_out with slippage (higher for sell)
            slippage = self.slippage_bps / 10000.0  # Convert from basis points to percentage
            min_out_lamports = 0  # For selling, we set min_out to 0 for speed
            min_out_bytes = min_out_lamports.to_bytes(8, byteorder="little")
            
            # Set amount_in (bytes 1-8)
            for i in range(8):
                instr_data[i+1] = amount_bytes[i]
                
            # Set min_out (bytes 9-16)
            for i in range(8):
                instr_data[i+9] = min_out_bytes[i]
            
            # Execute sell with Rust trading core
            tx_signature = self.trader.swap_with_accounts(
                instr_data,  # Proper 0x09 instruction data
                accounts.get("amm_id", ""),
                accounts.get("amm_authority", ""),
                accounts.get("amm_open_orders", ""),
                accounts.get("amm_target_orders", ""),
                accounts.get("pool_coin_token_account", ""),
                accounts.get("pool_pc_token_account", ""),
                accounts.get("serum_market", ""),
                accounts.get("serum_bids", ""),
                accounts.get("serum_asks", ""),
                accounts.get("serum_event_queue", ""),
                accounts.get("serum_coin_vault", ""),
                accounts.get("serum_pc_vault", ""),
                accounts.get("serum_vault_signer", ""),
                token_mint,
                False  # Is sell
            )
            
            return tx_signature
        
        except Exception as e:
            logger.error(f"Error selling tokens: {e}")
            return None
    
    async def execute_emergency_sell(self, opportunity):
        """Emergency sell to recover funds"""
        try:
            token_mint = opportunity.token_mint
            
            # Get current token balance
            token_balance = self.trader.get_token_balance(token_mint)
            
            if token_balance <= 0:
                logger.warning(f"No tokens to sell: {token_mint[:8]}")
                self.active_opportunity = None
                return
                
            logger.info(f"Emergency selling {token_balance} tokens of {token_mint[:8]}")
            
            # Get token decimals if not already known
            if opportunity.token_decimals == 9:  # Default value
                opportunity.token_decimals = self.trader.get_token_decimals(token_mint)
            
            # Execute sell with higher slippage
            for attempt in range(3):
                multiplier = 1.0 + (attempt * 0.5)  # Increase slippage with each attempt
                
                tx_signature = self.execute_token_sell(
                    token_mint=token_mint,
                    token_amount=token_balance,
                    token_decimals=opportunity.token_decimals,
                    accounts=opportunity.accounts
                )
                
                if tx_signature:
                    logger.info(f"Emergency sell successful on attempt {attempt+1}: {tx_signature[:8]}")
                    
                    # Get updated balances
                    async with self.balance_lock:
                        self.sol_balance = self.trader.get_sol_balance()
                        self.wsol_balance = self.trader.get_token_balance(
                            os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                        )
                        self.last_balance_update = time.time()
                        
                        # Calculate profit/loss
                        total_end = self.sol_balance + self.wsol_balance
                        profit = total_end - (opportunity.start_sol_balance + opportunity.start_wsol_balance)
                        
                        opportunity.profit = profit
                        opportunity.end_sol_balance = self.sol_balance
                        opportunity.end_wsol_balance = self.wsol_balance
                        
                        if profit >= 0:
                            self.stats.total_profit += profit
                            logger.info(f"Emergency sell profit: +{profit:.6f} SOL")
                        else:
                            self.stats.total_loss += abs(profit)
                            logger.info(f"Emergency sell loss: {profit:.6f} SOL")
                    
                    break
                else:
                    logger.error(f"Emergency sell attempt {attempt+1} failed")
                    await asyncio.sleep(0.5 * (attempt + 1))
            
            # Clear active opportunity
            self.active_opportunity = None
                
        except Exception as e:
            logger.error(f"Error in emergency sell: {e}")
            logger.error(traceback.format_exc())
            self.active_opportunity = None
    
    def estimate_profit(self, front_amount, victim_amount):
        """Estimate profit with enhanced calculation"""
        # Raydium fee
        raydium_fee = 0.0025  # 0.25%
        
        # Calculate token flow
        tokens_received = front_amount * (1.0 - raydium_fee)
        
        # Calculate price impact based on victim size
        # Larger victims cause more price impact
        price_impact_percent = min(0.05 * (victim_amount / 5.0) ** 0.5, 0.20)
        
        # Calculate expected SOL from back-run
        back_amount = tokens_received * (1.0 + price_impact_percent) * (1.0 - raydium_fee)
        
        # Calculate total fees
        priority_fee = (int(os.getenv("PRIORITY_FEE", "950000")) / LAMPORTS_PER_SOL) * (int(os.getenv("COMPUTE_LIMIT", "50000")) / 1000000) * 2
        total_fees = priority_fee + (0.000005 * 2)  # Base fee for two transactions
        
        # Calculate profit after fees
        raw_profit = back_amount - front_amount
        net_profit = raw_profit - total_fees
        
        # Apply safety margin
        adjusted_profit = net_profit * 0.95
        
        return {
            "raw_profit": raw_profit,
            "net_profit": net_profit,
            "adjusted_profit": adjusted_profit,
            "total_fees": total_fees,
            "price_impact": price_impact_percent,
            "is_profitable": adjusted_profit > self.min_profit
        }
    
    async def ensure_wsol(self, amount_needed):
        """Ensure sufficient WSOL for trading with reduced requirements"""
        # Already have enough
        if self.wsol_balance >= amount_needed:
            return True
            
        # Don't wrap too frequently
        if time.time() - self.last_wrap_time < self.wsol_wrap_interval:
            return False
            
        # Calculate amount to wrap (reduced buffer)
        wrap_amount = max(
            0.01,  # Minimum wrap amount reduced to 0.01
            (amount_needed - self.wsol_balance) * 1.1  # Buffer reduced from 1.5 to 1.1
        )
        
        # Ensure we have enough SOL - reduced min balance to improve success rate
        if self.sol_balance < wrap_amount + (self.min_sol_balance / 2):  # Reduced min balance requirement
            logger.error(f"Insufficient SOL: {self.sol_balance}, need {wrap_amount + (self.min_sol_balance / 2)}")
            return False
            
        try:
            logger.info(f"Wrapping {wrap_amount} SOL to WSOL")
            
            # Execute wrap
            signature = self.trader.wrap_sol(wrap_amount)
            self.last_wrap_time = time.time()
            
            # Wait for confirmation
            confirmed = self.trader.confirm_transaction(signature, 15)
            
            if confirmed:
                # Update balances
                self.sol_balance = self.trader.get_sol_balance()
                self.wsol_balance = self.trader.get_token_balance(
                    os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                )
                logger.info(f"Wrapped SOL successfully: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
                return True
            else:
                logger.error(f"Wrap confirmation failed")
                return False
                
        except Exception as e:
            logger.error(f"Error wrapping SOL: {e}")
            return False
    
    async def unwrap_wsol(self):
        """Unwrap WSOL back to native SOL"""
        try:
            logger.info(f"Unwrapping {self.wsol_balance} WSOL to SOL")
            
            # Execute unwrap
            signature = self.trader.unwrap_sol()
            
            # Wait for confirmation
            confirmed = self.trader.confirm_transaction(signature, 15)
            
            if confirmed:
                # Update balances
                self.sol_balance = self.trader.get_sol_balance()
                self.wsol_balance = self.trader.get_token_balance(
                    os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                )
                logger.info(f"Unwrapped WSOL successfully: SOL={self.sol_balance:.4f}, WSOL={self.wsol_balance:.4f}")
                return True
            else:
                logger.error(f"Unwrap confirmation failed")
                return False
                
        except Exception as e:
            logger.error(f"Error unwrapping WSOL: {e}")
            return False
    
    async def monitor_balances(self):
        """Monitor and manage SOL/WSOL balances"""
        check_interval = 60  # Check every minute
        
        while self.running:
            try:
                # Skip if active opportunity
                if self.active_opportunity:
                    await asyncio.sleep(check_interval)
                    continue
                
                # Get current balances
                current_time = time.time()
                if current_time - self.last_balance_update > 60:
                    self.sol_balance = self.trader.get_sol_balance()
                    self.wsol_balance = self.trader.get_token_balance(
                        os.getenv("WSOL_ADDRESS", "So11111111111111111111111111111111111111112")
                    )
                    self.last_balance_update = current_time
                
                # Check if we need to wrap SOL or unwrap WSOL
                total_balance = self.sol_balance + self.wsol_balance
                wsol_ratio = self.wsol_balance / total_balance if total_balance > 0 else 0
                target_ratio = float(os.getenv("WSOL_TARGET_RATIO", "0.75"))
                
                # If SOL balance is too low, unwrap WSOL
                if self.sol_balance < self.min_sol_balance and self.wsol_balance > self.min_sol_balance * 2:
                    logger.info(f"SOL balance too low ({self.sol_balance:.4f}), unwrapping WSOL")
                    await self.unwrap_wsol()
                
                # If WSOL ratio is too low, wrap SOL
                elif wsol_ratio < target_ratio * 0.8 and self.sol_balance > self.min_sol_balance * 1.5:
                    wrap_amount = (total_balance * target_ratio - self.wsol_balance) * 1.1
                    if wrap_amount > 0.01:  # Reduced min wrap amount
                        logger.info(f"WSOL ratio too low ({wsol_ratio:.2f}), wrapping {wrap_amount:.4f} SOL")
                        await self.ensure_wsol(wrap_amount)
                
            except Exception as e:
                logger.error(f"Error in balance monitoring: {e}")
                
            await asyncio.sleep(check_interval)

async def main():
    """Main entry point"""
    try:
        # Initialize and run the MEV trader
        trader = RaydiumMEVTrader()
        await trader.run()
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received, shutting down gracefully...")
    except Exception as e:
        logger.critical(f"Critical error: {e}")
        logger.critical(traceback.format_exc())

if __name__ == "__main__":
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())