# src/python/raydium_monitor.py

import asyncio
import logging
import websockets
from dataclasses import dataclass, field
from typing import Dict, Optional, Set, List, Deque
from decimal import Decimal
import time
import torch
import numpy as np
from collections import deque
from concurrent.futures import ThreadPoolExecutor
import base64
import json
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s.%(msecs)03d [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

@dataclass
class PoolState:
    base_reserve: Decimal = field(default_factory=lambda: Decimal('0'))
    quote_reserve: Decimal = field(default_factory=lambda: Decimal('0'))
    timestamp: float = field(default_factory=time.time)
    last_trade_price: float = 0.0
    trade_count: int = 0
    update_count: int = 0
    volatility: float = 0.0
    recent_trades: Deque[Dict] = field(default_factory=lambda: deque(maxlen=1000))
    volume_24h: float = 0.0
    price_impact: float = 0.0
    liquidity_depth: float = 0.0
    sandwich_opportunities: int = 0
    impact_opportunities: int = 0

@dataclass
class MonitorMetrics:
    messages_processed: int = 0
    opportunities_found: int = 0
    processing_time: float = 0.0
    successful_detections: int = 0
    failed_detections: int = 0
    peak_memory: float = 0.0
    start_time: float = field(default_factory=time.time)
    latency_stats: List[float] = field(default_factory=list)

class RaydiumMonitor:
    def __init__(self):
        """Initialize monitor with production optimizations"""
        self.active = False
        self.metrics = MonitorMetrics()
        
        # Load configuration from environment variables
        self.RAYDIUM_PROGRAM = os.getenv("RAYDIUM_PROGRAM_ID")
        self.HELIUS_WS_URL = os.getenv("HELIUS_WS_URL")
        
        if not self.RAYDIUM_PROGRAM or not self.HELIUS_WS_URL:
            raise ValueError("Missing required environment variables (RAYDIUM_PROGRAM_ID or HELIUS_WS_URL).")
        
        # Initialize GPU for calculations if available
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            
            # Dedicated CUDA streams
            self.streams = {
                'monitor': torch.cuda.Stream(priority=-2),  # Highest priority
                'process': torch.cuda.Stream(priority=-1),  # High priority
                'compute': torch.cuda.Stream(priority=0)    # Normal priority
            }
        
        # Pool state tracking
        self.pool_states: Dict[str, PoolState] = {}
        self.active_pools: Set[str] = set()
        self.last_update: Dict[str, float] = {}
        
        # Price tracking
        self.price_cache: Dict[str, deque] = {}
        self.price_trends: Dict[str, List[float]] = {}
        
        # Opportunity tracking
        self.opportunity_cache: Dict[str, Dict] = {}
        self.recent_victims: Dict[str, float] = {}
        self.execution_stats: Dict[str, List[float]] = {
            'latency': [],
            'slippage': [],
            'profit': []
        }
        
        # Initialize processing buffers
        self._setup_buffers()
        
        # Configure thread pool
        self.thread_pool = ThreadPoolExecutor(max_workers=2)
        
        logger.info(f"Enhanced price monitor initialized on {torch.cuda.get_device_name()}")

    def _setup_buffers(self):
        """Initialize GPU buffers with pinned memory"""
        try:
            with torch.cuda.stream(self.streams['compute']):
                # Price calculation buffer
                self.price_buffer = torch.zeros(
                    (1024, 64),
                    device=self.device,
                    dtype=torch.float16,
                    pin_memory=True
                ).contiguous()
                
                # Impact calculation buffer
                self.impact_buffer = torch.zeros(
                    (512, 32),
                    device=self.device,
                    dtype=torch.float32,
                    pin_memory=True
                ).contiguous()
                
                # Opportunity detection buffer
                self.detect_buffer = torch.zeros(
                    (256, 128),
                    device=self.device,
                    dtype=torch.float16,
                    pin_memory=True
                ).contiguous()
                
                torch.cuda.synchronize()
                logger.debug("GPU buffers initialized")
                
        except Exception as e:
            logger.error(f"Buffer initialization error: {e}")
            raise

    async def start_monitoring(self):
        """Start price monitoring with enhanced error handling"""
        self.active = True
        logger.info("Starting price monitoring...")
        
        # Initialize CUDA resources
        if not torch.cuda.is_available():
            logger.error("CUDA initialization failed")
            return
        
        while self.active:
            try:
                async with websockets.connect(
                    self.HELIUS_WS_URL,
                    ping_interval=30,
                    ping_timeout=10,
                    close_timeout=60,
                    max_size=16 * 1024 * 1024,
                    compression=None
                ) as ws:
                    await self._subscribe(ws)
                    await self._process_messages(ws)
                    
            except websockets.exceptions.ConnectionClosed:
                logger.warning("WebSocket connection closed. Reconnecting...")
                await asyncio.sleep(0.1)
            except Exception as e:
                logger.error(f"Monitoring error: {e}")
                await asyncio.sleep(0.1)

    async def _subscribe(self, ws):
        """Subscribe to relevant program and account updates"""
        try:
            # Subscribe to Raydium program
            await ws.send(json.dumps({
                "jsonrpc": "2.0",
                "id": "raydium_sub",
                "method": "programSubscribe",
                "params": [
                    self.RAYDIUM_PROGRAM,
                    {"encoding": "base64", "commitment": "processed"}
                ]
            }))
            
            # Subscribe to relevant market accounts
            accounts = [
                os.getenv('SOL_WSOL_AMM_ID'),
                os.getenv('SERUM_MARKET_ACCOUNT')
            ]
            
            for account in accounts:
                await ws.send(json.dumps({
                    "jsonrpc": "2.0",
                    "id": f"account_{account}",
                    "method": "accountSubscribe",
                    "params": [
                        account,
                        {"encoding": "base64", "commitment": "processed"}
                    ]
                }))
            
            logger.info("Subscriptions initialized")
            
        except Exception as e:
            logger.error(f"Subscription error: {e}")
            raise

    async def _process_messages(self, ws):
        """Process messages with batching and GPU acceleration"""
        while self.active:
            try:
                # Collect message batch
                messages = []
                batch_start = time.perf_counter()
                
                while len(messages) < 100:
                    try:
                        message = await asyncio.wait_for(ws.recv(), timeout=0.05)
                        messages.append(json.loads(message))
                    except asyncio.TimeoutError:
                        break
                    except Exception as e:
                        logger.error(f"Message receive error: {e}")
                        continue
                
                if not messages:
                    continue
                
                # Process batch with GPU acceleration
                if torch.cuda.is_available():
                    with torch.cuda.stream(self.streams['process']):
                        results = await self._process_message_batch(messages)
                else:
                    results = await self._process_message_batch(messages)
                
                # Update metrics
                batch_time = time.perf_counter() - batch_start
                self.metrics.processing_time = batch_time
                self.metrics.messages_processed += len(messages)
                
                # Log performance metrics
                if self.metrics.messages_processed % 1000 == 0:
                    self._log_metrics()
                
            except Exception as e:
                logger.error(f"Message processing error: {e}")
                await asyncio.sleep(0.1)

    async def _process_message_batch(self, messages: List[Dict]) -> List[Dict]:
        """Process message batch"""
        results = []

        try:
            for msg in messages:
                if 'params' not in msg:
                    continue
                
                params = msg['params']
                if 'result' not in params:
                    continue
                
                result = params['result']
                if 'value' not in result:
                    continue
                
                # Process Raydium program logs
                if "ray_log:" in str(result):
                    opportunity = await self._process_raydium_log(result)
                    if opportunity:
                        results.append(opportunity)
                        logger.info(f"Opportunity found: {opportunity}")
                
                # Process account updates
                elif 'data' in result.get('value', {}):
                    await self._process_account_update(result)
            
            return results
            
        except Exception as e:
            logger.error(f"Batch processing error: {e}")
            return []

    async def _process_raydium_log(self, log_data: Dict) -> Optional[Dict]:
        """Process Raydium log"""
        try:
            log_str = log_data.get('value', {}).get('logs', [''])[0]
            if not log_str.startswith('ray_log:'):
                return None
            
            raw_data = base64.b64decode(log_str.split('ray_log: ')[1].strip())
            
            amount_in = int.from_bytes(raw_data[0:8], 'little') / 1_000_000_000
            amount_out = int.from_bytes(raw_data[8:16], 'little') / 1_000_000_000
            pool_address = raw_data[64:96].hex()
            
            if amount_in < 0.1:  # Minimum amount check
                return None
            
            opportunity = await self.analyze_mempool_tx(raw_data, self.last_price)
            return opportunity
            
        except Exception as e:
            logger.error(f"Raydium log processing error: {e}")
            return None

    async def _process_account_update(self, update: Dict):
        """Process account update data"""
        try:
            account = update.get('pubkey')
            if not account:
                return
            
            data = update.get('value', {}).get('data')
            if not data:
                return
            
            if account == os.getenv('SOL_WSOL_AMM_ID'):
                await self._update_amm_state(data[0], account)
            elif account == os.getenv('SERUM_MARKET_ACCOUNT'):
                await self._update_market_state(data[0], account)
                
        except Exception as e:
            logger.error(f"Account update error: {e}")

    def _log_metrics(self):
        """Log performance metrics"""
        try:
            uptime = time.time() - self.metrics.start_time
            messages_per_second = self.metrics.messages_processed / uptime
            
            logger.info(f"""
Performance Metrics:
Messages Processed: {self.metrics.messages_processed}
Messages/Second: {messages_per_second:.2f}
Opportunities Found: {self.metrics.opportunities_found}
Success Rate: {(self.metrics.successful_detections / max(1, self.metrics.messages_processed)) * 100:.2f}%
Average Latency: {sum(self.metrics.latency_stats[-100:]) / max(1, len(self.metrics.latency_stats[-100:])) * 1000:.2f}ms
Peak Memory: {self.metrics.peak_memory:.2f}MB
            """)
            
        except Exception as e:
            logger.error(f"Metrics logging error: {e}")

    async def cleanup(self):
        """Cleanup resources properly"""
        logger.info("Cleaning up trade monitor...")
        try:
            self.active = False
            
            # Clean GPU resources
            if torch.cuda.is_available():
                torch.cuda.synchronize()
                for stream in self.streams.values():
                    stream.synchronize()
                torch.cuda.empty_cache()
            
            logger.info("Trade monitor cleanup completed")
            
        except Exception as e:
            logger.error(f"Error during cleanup: {e}")

 please fix this with  class RaydiumMonitor:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self._setup_buffers()
        
    def _setup_buffers(self):
        # Create buffers on CPU first, then move to GPU if available
        self.price_buffer = torch.zeros(
            (self.BUFFER_SIZE, self.PRICE_DIMENSIONS), 
            dtype=torch.float32
        )
        
        if self.device.type == 'cuda':
            self.price_buffer = self.price_buffer.cuda()