main.py

# Benchmark Avançado - Multithread + CPU + RAM

import platform
import psutil
import time
import os
import asyncio
import websockets
import threading
import numpy as np
from io import StringIO
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed


def print_system_info():
    print("=" * 50)
    print("INFORMAÇÕES DO SISTEMA")
    print("=" * 50)
    print(f"Plataforma:          {platform.system()} {platform.release()}")
    print(f"Processador:         {platform.processor()}")
    print(f"Núcleos Físicos:     {psutil.cpu_count(logical=False)}")
    print(f"Threads Lógicas:     {psutil.cpu_count(logical=True)}")
    print(f"Memória Total:       {psutil.virtual_memory().total / (1024 ** 3):.2f} GB")
    print(f"Memória Disponível:  {psutil.virtual_memory().available / (1024 ** 3):.2f} GB")
    print(f"Uso da CPU:          {psutil.cpu_percent(interval=1)}%")
    print(f"Uso da Memória:      {psutil.virtual_memory().percent}%")


# ── CPU ──────────────────────────────────────────────────────────────────────

def cpu_worker(thread_id: int, iteracoes: int) -> dict:
    """Tarefa de CPU executada por cada thread."""
    inicio = time.time()
    total = sum(i ** 2 for i in range(1, iteracoes))
    duracao = time.time() - inicio
    return {"thread_id": thread_id, "total": total, "tempo": duracao}


def cpu_benchmark_multithread(num_threads: int | None = None, iteracoes: int = 5_000_000) -> dict:
    num_threads = num_threads or psutil.cpu_count(logical=True)
    print(f"\n{'=' * 50}")
    print(f"BENCHMARK DE CPU — {num_threads} threads ({iteracoes:,} iterações cada)")
    print("=" * 50)

    resultados = []
    inicio_global = time.time()

    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = {
            executor.submit(cpu_worker, tid, iteracoes): tid
            for tid in range(num_threads)
        }
        for future in as_completed(futures):
            res = future.result()
            resultados.append(res)
            print(f"  Thread {res['thread_id']:>2}: {res['tempo']:.2f}s")

    tempo_total = time.time() - inicio_global
    print(f"\n  Tempo total (wall clock): {tempo_total:.2f}s")
    print(f"  Throughput:               {num_threads / tempo_total:.2f} threads/s")
    return {"threads": num_threads, "resultados": resultados, "tempo_total": tempo_total}


# ── RAM ──────────────────────────────────────────────────────────────────────

def ram_worker(thread_id: int, tamanho_mb: int) -> dict:
    """Aloca, escreve e lê um bloco de RAM; retorna largura de banda."""
    tamanho_bytes = tamanho_mb * 1024 * 1024
    elementos = tamanho_bytes // 8  # float64 = 8 bytes

    # Escrita
    inicio_escrita = time.time()
    dados = np.ones(elementos, dtype=np.float64)
    tempo_escrita = time.time() - inicio_escrita

    # Leitura
    inicio_leitura = time.time()
    _ = dados.sum()
    tempo_leitura = time.time() - inicio_leitura

    del dados

    bw_escrita = tamanho_mb / tempo_escrita if tempo_escrita > 0 else float("inf")
    bw_leitura = tamanho_mb / tempo_leitura if tempo_leitura > 0 else float("inf")

    return {
        "thread_id": thread_id,
        "tamanho_mb": tamanho_mb,
        "tempo_escrita": tempo_escrita,
        "tempo_leitura": tempo_leitura,
        "bw_escrita_mbs": bw_escrita,
        "bw_leitura_mbs": bw_leitura,
    }


def ram_benchmark_multithread(num_threads: int | None = None, tamanho_mb: int = 256) -> dict:
    num_threads = num_threads or psutil.cpu_count(logical=True)
    mem_disponivel_mb = psutil.virtual_memory().available / (1024 ** 2)
    tamanho_mb = min(tamanho_mb, int(mem_disponivel_mb * 0.8 / num_threads))

    print(f"\n{'=' * 50}")
    print(f"BENCHMARK DE RAM — {num_threads} threads ({tamanho_mb} MB cada)")
    print("=" * 50)

    resultados = []
    lock = threading.Lock()

    def executar(tid):
        res = ram_worker(tid, tamanho_mb)
        with lock:
            resultados.append(res)
            print(
                f"  Thread {tid:>2}: escrita {res['bw_escrita_mbs']:>8.1f} MB/s | "
                f"leitura {res['bw_leitura_mbs']:>8.1f} MB/s"
            )

    threads = [threading.Thread(target=executar, args=(i,)) for i in range(num_threads)]
    inicio_global = time.time()
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    tempo_total = time.time() - inicio_global

    bw_escrita_total = sum(r["bw_escrita_mbs"] for r in resultados)
    bw_leitura_total = sum(r["bw_leitura_mbs"] for r in resultados)
    print(f"\n  Largura de banda total — escrita: {bw_escrita_total:.1f} MB/s | leitura: {bw_leitura_total:.1f} MB/s")
    print(f"  Tempo total (wall clock): {tempo_total:.2f}s")

    return {
        "threads": num_threads,
        "tamanho_mb_por_thread": tamanho_mb,
        "resultados": resultados,
        "bw_escrita_total_mbs": bw_escrita_total,
        "bw_leitura_total_mbs": bw_leitura_total,
        "tempo_total": tempo_total,
    }


# ── WebSocket ─────────────────────────────────────────────────────────────────

async def enviar_resultado_ws(arquivo: str, ip: str = "10.2.129.12", porta: int = 8765):
    uri = f"ws://{ip}:{porta}"
    try:
        async with websockets.connect(uri) as ws:
            await ws.send(os.path.basename(arquivo))
            with open(arquivo, "r") as f:
                await ws.send(f.read())
        print(f"\nArquivo '{arquivo}' enviado para {ip}:{porta} via WebSocket.")
    except Exception as e:
        print(f"\nFalha ao enviar arquivo via WebSocket: {e}")


# ── Main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    NUM_THREADS = psutil.cpu_count(logical=True)   # ajuste se quiser fixar
    CPU_ITERACOES = 5_000_000                       # iterações por thread
    RAM_MB_POR_THREAD = 256                         # MB alocados por thread

    # Redireciona stdout para capturar tudo
    saida = StringIO()
    sys.stdout = saida

    print_system_info()
    resultado_cpu = cpu_benchmark_multithread(NUM_THREADS, CPU_ITERACOES)
    resultado_ram = ram_benchmark_multithread(NUM_THREADS, RAM_MB_POR_THREAD)

    sys.stdout = sys.__stdout__
    texto = saida.getvalue()
    print(texto)  # exibe no terminal também

    nome_arquivo = f"resultado_benchmark_{platform.node()}.txt"
    with open(nome_arquivo, "w") as f:
        f.write(texto)
    print(f"Resultado salvo em '{nome_arquivo}'")

    asyncio.run(enviar_resultado_ws(nome_arquivo))