#!/usr/bin/env python3
import socket
import sys
import time
import shlex

PORT = 5056
ALL_IPS = [
    "172.22.156.3",
    "172.22.158.2",
    "172.22.94.2",
    "172.22.156.4",
    "172.22.158.3",
    "172.22.94.3",
    "172.22.156.5",
    "172.22.158.4",
    "172.22.94.4",
    "172.22.156.6"
]

def send_message(sock, message):
    e_data = message.encode('utf-8')
    #header for the message length
    header = str(len(e_data)).zfill(10).encode('utf-8')
    sock.sendall(header)
    sock.sendall(e_data)

def receive_exact(sock, num_bytes):
    received_data = b""
    while len(received_data) < num_bytes:
        chunk = sock.recv(num_bytes - len(received_data))
        if not chunk:
            raise ConnectionError("connection closed before the full message arrived")
        received_data += chunk
    return received_data

def receive_message(sock):
    header = receive_exact(sock, 10)
    if not header:
        return None
    message_length = int(header.decode('utf-8'))
    received_data = receive_exact(sock, message_length)
    return received_data.decode('utf-8')

def grep_all_ips(ips, cmd, if_unit_test):
    match_count = []

    for i, ip in enumerate(ips):
        try:
            client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            client_socket.connect((ip, PORT))

            filename = f"machine.{i + 1}.test.log" if if_unit_test else f"machine.{i + 1}.log"

            full_cmd = shlex.join(cmd + [filename])

            send_message(client_socket, full_cmd)
            data = receive_message(client_socket)

            count = len(data.splitlines())
            match_count.append(count)
            if count > 0:
                print(f"Filename: machine.{i + 1}.log")
                print(f"Matches: {count}")
                print(data)

        except Exception as e:
            print(f"Failed to connect to {ip}:{PORT} - {e}")
            match_count.append(-1)
        finally:
            client_socket.close()
    
    return match_count

#main
#if unit test flag exists, then use the test log files instead of the actual log files
cmd = ["grep", *[arg for arg in sys.argv[1:] if arg != "-t"]]
#track time
start_time = time.perf_counter()
is_unit_test = "-t" in sys.argv
match_count = grep_all_ips(ALL_IPS, cmd, if_unit_test=is_unit_test)
end_time = time.perf_counter()

print("=========================== SUMMARY ===========================")
for i, count in enumerate(match_count):
    if count == -1:
        if [x for x in sys.argv].count("-t") > 0:
            print(f"machine.{i + 1}.test.log: Connection failed")
        else:   
            print(f"machine.{i + 1}.log: Connection failed")
    else:
        if [x for x in sys.argv].count("-t") > 0:
            print(f"machine.{i + 1}.test.log: {count} matches")
        else:
            print(f"machine.{i + 1}.log: {count} matches")
total_matches = sum(count for count in match_count if count > 0)
print(f"Total matches: {total_matches}")
print("===============================================================")
print(f"Time taken: {end_time - start_time:.2f} seconds")