import asyncio
import json 
import os 
from datetime import datetime
import pytz 
from websockets import connect 
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich.align import Align
from rich.layout import Layout
from rich.live import Live
from rich import box
from rich.progress import Progress, BarColumn, TextColumn, SpinnerColumn
from rich.columns import Columns
from rich.rule import Rule
from collections import defaultdict, deque
import time
import statistics

# Initialize Rich console
console = Console()

# List of symbols you want to track 
symbols = ['btcusdt', 'ethusdt', 'solusdt', 'bnbusdt', 'dogeusdt', 'wifusdt']
websocket_url_base = 'wss://fstream.binance.com/ws/'

def get_trade_style(usd_size, trade_type):
    """Get enhanced styling for trade based on size and type"""
    if usd_size >= 1000000:  # $1M+
        return {
            'emoji': '💎',
            'color': 'bold bright_magenta' if trade_type == 'SELL' else 'bold bright_cyan',
            'border': 'heavy',
            'title': '🌊 MEGA WHALE DETECTED 🌊',
            'bg_color': 'on_blue' if trade_type == 'BUY' else 'on_red'
        }
    elif usd_size >= 500000:  # $500K+
        return {
            'emoji': '🏦',
            'color': 'bold magenta' if trade_type == 'SELL' else 'bold cyan',
            'border': 'double',
            'title': '🔥 MASSIVE WHALE 🔥',
            'bg_color': None
        }
    elif usd_size >= 100000:  # $100K+
        return {
            'emoji': '💰',
            'color': 'bold red' if trade_type == 'SELL' else 'bold green',
            'border': 'rounded',
            'title': '⚡️ WHALE ALERT ⚡️',
            'bg_color': None
        }
    else:  # $15K+
        return {
            'emoji': '💵',
            'color': 'red' if trade_type == 'SELL' else 'green',
            'border': 'ascii',
            'title': None,
            'bg_color': None
        }


async def binance_trade_stream(url, symbol):
    """Enhanced WebSocket connection for individual symbol"""
    async with connect(url) as websocket:
        while True:
            try:
                message = await websocket.recv()
                data = json.loads(message)
                
                # Parse trade data
                price = float(data['p'])
                quantity = float(data['q'])
                trade_time = int(data['T'])
                is_buyer_maker = data['m']
                
                # Calculate USD size
                usd_size = price * quantity 
                
                if usd_size >= 15000:
                    # Determine trade type
                    trade_type = 'SELL' if is_buyer_maker else "BUY"
                 
                    # Format time
                    est = pytz.timezone('US/Eastern')
                    readable_trade_time = datetime.fromtimestamp(trade_time / 1000, est).strftime('%H:%M:%S')
                    
                    display_symbol = symbol.upper().replace('USDT', '')
                    
                    # Get styling
                    style_info = get_trade_style(usd_size, trade_type)
                    
                    # Create enhanced trade notification
                    trade_text = Text()
                    trade_text.append(f"[{readable_trade_time}] ", style="dim bright_white")
                    trade_text.append(f"{style_info['emoji']} ", style="bold")
                    trade_text.append(f"{trade_type:<4} ", style=f"bold {'bright_green' if trade_type == 'BUY' else 'bright_red'}")
                    trade_text.append(f"{display_symbol:<4} ", style="bold bright_cyan")
                    trade_text.append(f"${usd_size:,.0f} ", style="bold bright_yellow")
                    trade_text.append(f"@ ${price:,.2f} ", style="bright_white")

                    # Create panel for large trades (100K+)
                    if usd_size >= 100000 and style_info['title']:
                        enhanced_content = Text()
                        enhanced_content.append(f"{style_info['emoji']} TRADE SIZE: ${usd_size:,.0f}\n", style="bold bright_yellow")
                        enhanced_content.append(f"💲 PRICE: ${price:,.2f}\n", style="bright_white")
                        enhanced_content.append(f"📊 SYMBOL: {display_symbol}\n", style="bold bright_cyan")
                        enhanced_content.append(f"⏰ TIME: {readable_trade_time}", style="dim")
                        
                        panel = Panel(
                            Align.center(enhanced_content),
                            title=style_info['title'],
                            subtitle=f"🐋 {trade_type} ORDER",
                            box=getattr(box, style_info['border'].upper()),
                            border_style=style_info['color'],
                            padding=(1, 2)
                        )
                        console.print(panel)
                        
                        # Add separator for mega trades
                        if usd_size >= 1000000:
                            console.print(Rule("🌊🌊🌊", style="bold bright_blue"))
                    else:
                        console.print(trade_text)

            except Exception as e:
                console.print(f"[red]❌ Error in {symbol}: {e}[/red]")
                await asyncio.sleep(5)

async def dashboard_updater():
    """Periodically update the dashboard"""
    while True:
        try:
            await asyncio.sleep(1)  # Update every 10 seconds
            # You could implement a live dashboard here using Rich Live
            # For now, we'll just continue with the stream-based approach
        except Exception as e:
            console.print(f"[red]Dashboard error: {e}[/red]")

async def main():
    """Enhanced main function"""
    console.clear()
    
    # Enhanced startup sequence
    startup_progress = Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
    )
    
    with startup_progress:
        task = startup_progress.add_task("🐋 Initializing Whale Monitor...", total=100)
        
        for i in range(100):
            await asyncio.sleep(0.02)
            startup_progress.update(task, advance=1)
            
            if i == 20:
                startup_progress.update(task, description="🌐 Connecting to Binance...")
            elif i == 50:
                startup_progress.update(task, description="📊 Setting up data streams...")
            elif i == 80:
                startup_progress.update(task, description="🎯 Calibrating whale detection...")
    
    console.clear()
    
    # Display enhanced startup message
    startup_panel = Panel(
        Align.center(
            Text("🐋 BINANCE WHALE MONITOR ACTIVE 🐋\n\n") +
            Text("💎 Tracking: ", style="bright_green") +
            Text(", ".join([s.upper().replace('USDT', '') for s in symbols]), style="bright_yellow") +
            Text("\n🎯 Minimum Trade Size: $15,000", style="bright_white") +
            Text("\n⚡️ Real-time WebSocket Streams Connected", style="bright_cyan") +
            Text("\n\n🚀 Ready to catch whales!", style="bold bright_magenta")
        ),
        title="🌊 WHALE MONITOR INITIALIZED 🌊",
        border_style="bright_green",
        box=box.DOUBLE,
        padding=(1, 2)
    )
    console.print(startup_panel)
    console.print()
    
    # Create tasks for each symbol trade stream 
    tasks = []
    
    # Add WebSocket tasks
    for symbol in symbols:
        stream_url = f"{websocket_url_base}{symbol}@aggTrade"
        tasks.append(binance_trade_stream(stream_url, symbol))
    
    # Add dashboard updater
    tasks.append(dashboard_updater())
    
    await asyncio.gather(*tasks)

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        console.print("\n")
        shutdown_panel = Panel(
            Align.center("🐋 Whale Monitor shutting down...\n👋 Thanks for whale watching!"),
            title="🌊 SHUTDOWN COMPLETE 🌊",
            border_style="bright_yellow",
            box=box.DOUBLE
        )
        console.print(shutdown_panel)
    except Exception as e:
        console.print(f"[red]❌ Critical Error: {e}[/red]")