#!/usr/bin/env python3
import socket, random, string, time, threading

# ————— Parámetros —————
NLB_HOST     = "nlb-recepcion-tcp-98af84506c9b42a6.elb.us-east-1.amazonaws.com"
TCP_PORT     = 7001
TOTAL_PKTS   = 10000     # total de paquetes a enviar
TARGET_PPS   = 100        # paquetes por segundo deseados
THREAD_COUNT = 10         # número de hilos concurrentes

# ————— Generador de mensajes —————
def generar_mensaje():
    equipo = ''.join(random.choices(string.digits, k=15))
    lon = round(random.uniform(-180,180),6)
    lat = round(random.uniform(-90,90),6)
    fh  = time.strftime("%Y%m%d%H%M%S")
    hx  = ''.join(random.choices(string.hexdigits.lower(), k=4))
    u4  = ''.join(random.choices(string.hexdigits.lower(), k=4))
    msg = (f"+RESP:GTFRI,500101,{equipo},GV55W,12703,10,1,1,0.0,229,36.3,"
           f"{lon},{lat},{fh},{hx},0007,1168,16C073,00,0.0,,,,96,110000,,,,"
           f"{fh},{u4}$")
    return msg.encode()

def worker(stop_event, shared, lock, pps_thread):
    if pps_thread <= 0:
        return
    interval = 1.0 / pps_thread
    next_send = time.perf_counter()

    while not stop_event.is_set():
        now = time.perf_counter()
        if now < next_send:
            time.sleep(next_send - now)
        next_send += interval

        # Reservar un envío
        with lock:
            if shared['sent'] >= TOTAL_PKTS:
                stop_event.set()
                break
            shared['sent'] += 1

        # Enviar en una conexión nueva
        try:
            sock = socket.create_connection((NLB_HOST, TCP_PORT), timeout=2)
            sock.sendall(generar_mensaje())
            sock.close()
        except Exception as e:
            # Si falla, lo ignoramos y seguimos; la cuenta ya fue descontada
            pass

def main():
    # Repartimos TARGET_PPS entre THREAD_COUNT hilos
    base, extra = divmod(TARGET_PPS, THREAD_COUNT)
    rates = [base + (1 if i < extra else 0) for i in range(THREAD_COUNT)]

    stop_event = threading.Event()
    lock = threading.Lock()
    shared = {'sent': 0}

    threads = []
    start = time.time()
    print(f"Iniciando prueba: {TOTAL_PKTS} paquetes a {TARGET_PPS} pkt/s usando {THREAD_COUNT} hilos")

    for pps in rates:
        t = threading.Thread(target=worker, args=(stop_event, shared, lock, pps), daemon=True)
        threads.append(t)
        t.start()

    # Esperamos que terminen
    for t in threads:
        t.join()

    end = time.time()
    duration = end - start
    sent = shared['sent']
    real_pps = sent / duration if duration > 0 else 0

    print("\n=== Resultados de la prueba ===")
    print(f"Paquetes esperados:          {TOTAL_PKTS}")
    print(f"Paquetes realmente enviados: {sent}")
    print(f"Duración:                    {duration:.2f} seg")
    print(f"Tasa promedio:               {real_pps:.2f} pkt/s")

if __name__ == "__main__":
    main()