SD/main/src/main/java/sd/protocol/SocketConnection.java

package sd.protocol; 

import java.io.Closeable;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;

import java.io.InputStream;
import java.io.OutputStream;
import java.net.ConnectException;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.util.concurrent.TimeUnit;

import sd.serialization.MessageSerializer;
import sd.serialization.SerializationException;
import sd.serialization.SerializerFactory;


/**
 * Wrapper class that simplifies communication via Sockets.
 * Includes connection retry logic for robustness.
 */
public class SocketConnection implements Closeable {

    private final Socket socket;
    private final OutputStream outputStream;
    private final InputStream inputStream;
    private final MessageSerializer serializer;

    // --- Configuration for Retry Logic ---
    /** Maximum number of connection attempts. */
    private static final int MAX_RETRIES = 5; 
    /** Delay between retry attempts in milliseconds. */
    private static final long RETRY_DELAY_MS = 1000; 

    /**
     * Constructor for the "Client" (who initiates the connection).
     * Tries to connect to a process that is already listening (Server).
     * Includes retry logic in case of initial connection failure.
     *
     * @param host The host address (e.g., "localhost" from your simulation.properties)
     * @param port The port (e.g., 8001 from your simulation.properties)
     * @throws IOException If connection fails after all retries.
     * @throws UnknownHostException If the host is not found (this error usually doesn't need retry).
     * @throws InterruptedException If the thread is interrupted while waiting between retries.
     */
    public SocketConnection(String host, int port) throws IOException, UnknownHostException, InterruptedException {
        Socket tempSocket = null;
        IOException lastException = null;

        System.out.printf("[SocketConnection] Attempting to connect to %s:%d...%n", host, port);

        // --- Retry Loop ---
        for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) {
            try {
                // Try to establish the connection
                tempSocket = new Socket(host, port);

                // If successful, break out of the retry loop
                System.out.printf("[SocketConnection] Connected successfully on attempt %d.%n", attempt);
                lastException = null; // Clear last error on success
                break;

            } catch (ConnectException | SocketTimeoutException e) {
                // These are common errors indicating the server might not be ready.
                lastException = e;
                System.out.printf("[SocketConnection] Attempt %d/%d failed: %s. Retrying in %d ms...%n",
                        attempt, MAX_RETRIES, e.getMessage(), RETRY_DELAY_MS);

                if (attempt < MAX_RETRIES) {
                    // Wait before the next attempt
                    TimeUnit.MILLISECONDS.sleep(RETRY_DELAY_MS);
                }
            } catch (IOException e) {
                // Other IOExceptions might be more permanent, but we retry anyway.
                 lastException = e;
                 System.out.printf("[SocketConnection] Attempt %d/%d failed with IOException: %s. Retrying in %d ms...%n",
                        attempt, MAX_RETRIES, e.getMessage(), RETRY_DELAY_MS);
                 if (attempt < MAX_RETRIES) {
                     TimeUnit.MILLISECONDS.sleep(RETRY_DELAY_MS);
                 }
            }
        } // --- End of Retry Loop ---

        // If after all retries tempSocket is still null, it means connection failed permanently.
        if (tempSocket == null) {
            System.err.printf("[SocketConnection] Failed to connect to %s:%d after %d attempts.%n", host, port, MAX_RETRIES);
            if (lastException != null) {
                 throw lastException; // Throw the last exception encountered
            } else {
                 // Should not happen if loop ran, but as a fallback
                 throw new IOException("Failed to connect after " + MAX_RETRIES + " attempts, reason unknown.");
            }
        }

        // If connection was successful, assign to final variable and create streams
        this.socket = tempSocket;
          
        this.outputStream = socket.getOutputStream();
        this.inputStream = socket.getInputStream();
        this.serializer = SerializerFactory.createDefault();

    }


    /**
     * Constructor for the "Server" (who accepts the connection).
     * Receives a Socket that has already been accepted by a ServerSocket.
     * No retry logic needed here as the connection is already established.
     *
     * @param acceptedSocket The Socket returned by serverSocket.accept().
     * @throws IOException If stream creation fails.
     */
    public SocketConnection(Socket acceptedSocket) throws IOException {
        this.socket = acceptedSocket;
        this.outputStream = socket.getOutputStream();
        this.inputStream = socket.getInputStream();
        this.serializer = SerializerFactory.createDefault();
        
    }

    /**
     * Sends (serializes) a MessageProtocol object over the socket.
     *
     * @param message The "envelope" (which contains the Vehicle) to be sent.
     * @throws IOException If writing to the stream fails or socket is not connected.
     */
    public void sendMessage(MessageProtocol message) throws IOException {
   if (socket == null || !socket.isConnected()) {
            throw new IOException("Socket is not connected");
        }

        try {
            // Serializa para bytes JSON
            byte[] data = serializer.serialize(message);
            
            // Write 4-byte length prefix
            DataOutputStream dataOut = new DataOutputStream(outputStream);
            dataOut.writeInt(data.length);
            dataOut.write(data);
            dataOut.flush();
            
        } catch (SerializationException e) {
            throw new IOException("Failed to serialize message", e);
        }
    }

    /**
     * Tries to read (deserialize) a MessageProtocol object from the socket.
     *
     * @return The "envelope" (MessageProtocol) that was received.
     * @throws IOException If the connection is lost, the stream is corrupted, or socket is not connected.
     * @throws ClassNotFoundException If the received object is unknown.
     */
    public MessageProtocol receiveMessage() throws IOException, ClassNotFoundException {
       if (socket == null || !socket.isConnected()) {
            throw new IOException("Socket is not connected");
        }

        try {
            // Lê um prefixo de 4 bytes - indicador de tamanho
            DataInputStream dataIn = new DataInputStream(inputStream);
            int length = dataIn.readInt();
            
            if (length <= 0 || length > 10_000_000) { // Sanity check (10MB max)
                throw new IOException("Invalid message length: " + length);
            }
            
            // Ler dados da mensagem
            byte[] data = new byte[length];
            dataIn.readFully(data);
            
            // Deserialize do JSON
            return serializer.deserialize(data, MessageProtocol.class);
            
        } catch (SerializationException e) {
            throw new IOException("Failed to deserialize message", e);
        }
    }

    /**
     * Closes the socket and all streams (Input and Output).
     */
    @Override
    public void close() throws IOException {
        if (inputStream != null) inputStream.close();
        if (outputStream != null) outputStream.close();
        if (socket != null) socket.close();
    }

    /**
     * @return true if the socket is still connected and not closed.
     */
    public boolean isConnected() {
        return socket != null && socket.isConnected() && !socket.isClosed();
    }
}