import asyncio
import json 
from datetime import datetime 
from websockets import connect 
import random 

from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.live import Live
from rich.layout import Layout
from rich.text import Text
from rich.align import Align
from rich import box


console = Console()
symbols = ['btcusdt', 'ethusdt', 'solusdt', 'bnbusdt', 'dogeusdt', 'wifusdt']
websocket_url_base = 'wss://fstream.binance.com/ws/'

# Store latest data for each symbol
funding_data = {symbol: None for symbol in symbols}
print_lock = asyncio.Lock()

def get_funding_style(yearly_rate):
    """Return style based on funding rate"""
    if yearly_rate > 50:
        return "bold white on red"
    elif yearly_rate > 30:
        return "bold black on yellow"
    elif yearly_rate > 5:
        return "bold white on blue"
    elif yearly_rate < -10:
        return "bold white on green"
    else:
        return "bold black on bright_green"

def create_funding_table():
    """Create a rich table with current funding data"""
    table = Table(
        title="🚀 Binance Perpetual Funding Rates",
        title_style="bold cyan",
        box=box.ROUNDED,
        show_header=True,
        header_style="bold magenta"
    )
    
    table.add_column("Symbol", style="cyan", width=12)
    table.add_column("Funding Rate", justify="right", width=15)
    table.add_column("Yearly Rate", justify="right", width=15)
    table.add_column("Status", justify="center", width=20)
    table.add_column("Last Update", style="dim", width=12)
    
    for symbol in symbols:
        data = funding_data[symbol]
        if data:
            symbol_display = data['symbol_display']
            funding_rate = data['funding_rate']
            yearly_rate = data['yearly_rate']
            last_update = data['last_update']
            
            # Create status indicator
            if yearly_rate > 30:
                status = Text("🔥 HIGH", style="bold red")
            elif yearly_rate > 5:
                status = Text("📈 POSITIVE", style="bold yellow")
            elif yearly_rate < -10:
                status = Text("📉 NEGATIVE", style="bold green")
            else:
                status = Text("😐 NEUTRAL", style="dim")
            
            # Style the rates
            funding_text = Text(f"{funding_rate:.4f}%", style=get_funding_style(yearly_rate))
            yearly_text = Text(f"{yearly_rate:+.2f}%", style=get_funding_style(yearly_rate))
            
            table.add_row(
                symbol_display,
                funding_text,
                yearly_text,
                status,
                last_update
            )
        else:
            table.add_row(
                symbol.replace('usdt', '').upper(),
                Text("Loading...", style="dim"),
                Text("Loading...", style="dim"),
                Text("⏳ WAITING", style="dim"),
                Text("--:--:--", style="dim")
            )
    
    return table

def create_layout():
    """Create the main layout"""
    layout = Layout()
    
    # Header
    header = Panel(
        Align.center(
            Text("Binance Funding Rate Monitor", style="bold white"),
            vertical="middle"
        ),
        height=3,
        style="blue"
    )
    
    # Main table
    table = create_funding_table()
    
    # Footer with info
    footer_text = Text()
    footer_text.append("🟥 >50% ", style="bold white on red")
    footer_text.append("🟨 >30% ", style="bold black on yellow") 
    footer_text.append("🟦 >5% ", style="bold white on blue")
    footer_text.append("🟩 <-10% ", style="bold white on green")
    footer_text.append("💚 Normal", style="bold black on bright_green")
    
    footer = Panel(
        Align.center(footer_text),
        title="Legend",
        height=3,
        style="dim"
    )
    
    layout.split_column(
        Layout(header, name="header", size=3),
        Layout(table, name="main"),
        Layout(footer, name="footer", size=3)
    )
    
    return layout


async def binance_funding_stream(symbol):
    """Connect to Binance WebSocket and stream funding data"""
    global funding_data, print_lock
    
    websocket_url = f'{websocket_url_base}{symbol}@markPrice'
    
    while True:  # Reconnection loop
        try:
            async with connect(websocket_url) as websocket:
                while True:
                    try:
                        message = await websocket.recv()
                        data = json.loads(message)
                        
                        event_time = datetime.fromtimestamp(data['E'] / 1000).strftime("%H:%M:%S")
                        symbol_display = data['s'].replace('USDT', '').upper()
                        funding_rate = float(data['r']) * 100  # Convert to percentage
                        yearly_funding_rate = funding_rate * 3 * 365  # 3 times per day, 365 days
                        
                        async with print_lock:
                            funding_data[symbol] = {
                                'symbol_display': symbol_display,
                                'funding_rate': funding_rate,
                                'yearly_rate': yearly_funding_rate,
                                'last_update': event_time
                            }
                            
                    except json.JSONDecodeError:
                        continue
                    except Exception as e:
                        console.print(f"[red]Error processing data for {symbol}: {e}[/red]")
                        break
                        
        except Exception as e:
            console.print(f"[red]Connection error for {symbol}: {e}. Retrying in 5 seconds...[/red]")
            await asyncio.sleep(5)


async def display_loop():
    """Main display loop using Rich Live"""
    with Live(create_layout(), refresh_per_second=2, screen=True) as live:
        while True:
            live.update(create_layout())
            await asyncio.sleep(0.5)

async def main():
    """Main function to run all tasks"""
    console.print("[bold green]🚀 Starting Binance Funding Rate Monitor...[/bold green]")
    console.print("[yellow]Press Ctrl+C to exit[/yellow]\n")
    
    # Start WebSocket streams
    stream_tasks = [binance_funding_stream(symbol) for symbol in symbols]
    
    # Start display
    display_task = display_loop()
    
    # Run all tasks
    try:
        await asyncio.gather(*stream_tasks, display_task)
    except KeyboardInterrupt:
        console.print("\n[bold red]Shutting down...[/bold red]")


if __name__ == "__main__":
    asyncio.run(main())