import json
import asyncio
import sys, time
from quart import Quart, websocket

peer_list = {}
TRANSFERRED_DATA = 0

app = Quart(__name__)


def convert_bytes(byte_value):
    if byte_value < 1024:  # Less than 1 kilobyte
        return str(byte_value) +  " B"
    elif byte_value < 1024**2:  # Less than 1 megabyte
        kb_value = round(byte_value / 1024)
        return str(kb_value) + " KB"
    elif byte_value < 1024**3:  # Less than 1 gigabyte
        mb_value = round(byte_value / (1024**2),2)
        return str(mb_value) + " MB"
    elif byte_value < 1024**4:  # Less than 1 terabyte
        gb_value = round(byte_value / (1024**3),2)
        return str(gb_value) + " GB"
    elif byte_value < 1024**5:  # Less than 1 petabyte
        tb_value = round(byte_value / (1024**4),2)
        return str(tb_value) + " TB"
    else:  # More than or equal to 1 petabyte
        tb_value = round(byte_value / (1024**5),2)
        return str(tb_value) + " PB"

@app.route("/")
async def retmain():
    return f"Ready to relay\n"

@app.websocket("/ws")
async def handle_websockets():
    global TRANSFERRED_DATA
    peer_group_id = None
    while True:
        try:
            message = await websocket.receive()
            msg = json.loads(message)
            if "peer_group_id" in msg:
                peer_group_id = msg["peer_group_id"]
                if peer_group_id not in peer_list:
                    peer_list[peer_group_id] = set()
                peer_list[peer_group_id].add(websocket._get_current_object())
            for peer in peer_list[peer_group_id]:
                if peer != websocket._get_current_object():
                    await peer.send(message)
                    TRANSFERRED_DATA = TRANSFERRED_DATA + sys.getsizeof(message)
        except asyncio.exceptions.CancelledError:
            peer_list[peer_group_id].remove(websocket._get_current_object())
            if len(peer_list[peer_group_id]) < 1:
                peer_list.pop(peer_group_id)

@app.route("/stats")
async def return_stats():
    process_start_time = time.process_time()
    current_time = time.time()
    uptime = current_time - process_start_time
    uptime_str = time.strftime("%H:%M:%S", time.gmtime(uptime))
    peers_connected = 0
    for p in peer_list.values():
        peers_connected += len(p)
    resp = f"<p><b>Uptime: </b>{uptime_str}</p><p><b>Peers: </b>{peers_connected}</p><p><b>Transferred: </b>{convert_bytes(TRANSFERRED_DATA)}</p>"
    return resp

app.run()