#!/usr/bin/env python3

"""
Waddle command line interface to run either a server or a websocket client to write metrics to the console.

In server mode, we start the Waddle server and point to an existing database and/or folder to monitor for ingest.
In client mode, we can listen to a waddle server and stream metrics to the console.
"""

import asyncio
import websockets
import argparse
import json
from waddle import server as server_main
from rich.live import Live
from rich.table import Table
from rich.console import Console

async def websocket_client(uri):
    console = Console()
    log_entries = []

    def generate_table():
        table = Table(title="Waddle Metrics")
        table.add_column("Step", justify="right", style="cyan", no_wrap=True)
        table.add_column("Category", style="magenta")
        table.add_column("Name", style="green")
        table.add_column("Value", style="white")
        table.add_column("Timestamp", style="yellow")

        # Display the last 20 log entries
        for entry in log_entries[-20:]:
            table.add_row(
                str(entry.get('step', '')),
                entry.get('category', ''),
                entry.get('name', ''),
                str(entry.get('value', '')),
                entry.get('timestamp', '')
            )
        return table

    while True:
        try:
            async with websockets.connect(uri) as websocket:
                with Live(generate_table(), refresh_per_second=4, console=console) as live:
                    async for message in websocket:
                        log_entry = json.loads(message)
                        log_entries.append(log_entry)
                        live.update(generate_table())
        except Exception as e:
            console.print(f"[red]Connection error:[/] {e}")
            console.print("Reconnecting in 5 seconds...")
            await asyncio.sleep(5)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', type=str, choices=['server', 'client'], default='server')
    parser.add_argument('--log-level', type=str, default='critical')
    parser.add_argument('--server-port', type=int, default=8000)
    parser.add_argument('--server-bind', type=str, default='127.0.0.1')
    parser.add_argument('--server-url', type=str, default=None)
    parser.add_argument('--db-root', type=str, default='.waddle')
    parser.add_argument('--project', type=str, default='experiment')
    parser.add_argument('--watch-folder', type=str, default=None)
    args = parser.parse_args()

    if args.mode == 'server':
        # Start the Waddle server
        server_main.main(
            port=args.server_port,
            bind=args.server_bind,
            log_level=args.log_level,
        )
    elif args.mode == 'client':
        server_url = args.server_url or f'ws://{args.bind}:{args.server_port}/ws'
        asyncio.run(websocket_client(server_url))
