import asyncio
import logging
import sys
import time

from asyncio import StreamReader, StreamWriter

from cintes import genera_missatges

    
async def worker(id_worker: int, host: str, port: int, stats: dict, end_time: float):
    """
    Simula un canal de comunicació constant que envia dades de telemetria
    recurrentment fins que s'acaba el temps de la prova.
    """
    g = genera_missatges(f"cinta_{id_worker:02d}")
    try:
        # Obrim la connexió de stream
        reader, writer = await asyncio.open_connection(host, port)

        while time.monotonic() < end_time:
            try:

                # Enviem una dada de telemetria (alternant ALERTA per provar lògica)
                missatge = next(g)

                writer.write(missatge.encode())
                await writer.drain()  # Garantim que les dades surten cap al servidor

                # Esperem la resposta del servidor (protocol de telemetria)
                r = await reader.readline()

                stats['requests'] += 1

            except Exception:
                stats['errors'] += 1
    except Exception:
        logging.exception(f"[Worker {id_worker:2d}] ERROR, no es pot establir la connexió")
        stats['errors'] += 1
    finally:
        # Tanquem la connexió de manera neta
        writer.close()
        await writer.wait_closed()


async def run_prova_de_carrega(host: str, port: int, connections: int, duration: int):
    stats = {'requests': 0, 'errors': 0}
    start_time = time.monotonic()
    end_time = start_time + duration
    
    # Creem el grup de tasques que mantindran la concurrència
    tasks = [
        asyncio.create_task(worker(i, host, port, stats, end_time))
        for i in range(connections)
    ]
    
    # Esperem que tots els treballadors finalitzin quan s'acabi el temps
    await asyncio.gather(*tasks)
    
    total_time = time.monotonic() - start_time

    return stats, total_time


if __name__ == "__main__":
    # Paràmetres configurables de la prova
    HOST = '127.0.0.1'
    PORT = 7777
    CONCURRENCIA, DURADA_SEGONS = (int(a) for a in sys.argv[1:])
    
    try:
        print(f"Iniciant prova de {DURADA_SEGONS} segons amb {CONCURRENCIA} connexions concurrents...")
        print(f"Protocol: TCP Stream sobre {HOST}:{PORT}\n")
        
        stats, total_time = asyncio.run(run_prova_de_carrega(HOST, PORT, CONCURRENCIA, DURADA_SEGONS))
        rps = stats['requests'] / total_time
    
        print("-" * 45)
        print(f"Peticions totals: {stats['requests']}")
        print(f"Errors totals:    {stats['errors']}")
        print(f"Temps real:       {total_time:.2f} segons")
        print(f"Rendiment (RPS):  {rps:.2f} peticions/segon")
        print("-" * 45)
    except KeyboardInterrupt:
        pass