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 java.util.concurrent.locks.LockSupport;

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

/**
 * Wrapper de alto nível para gestão robusta de conexões TCP.
 * <p>
 * Esta classe abstrai a complexidade da API nativa {@link java.net.Socket}, oferecendo:
 * <ol>
 * <li><b>Resiliência:</b> Lógica de reconexão automática (Retry Loop) no arranque, crucial para sistemas
 * distribuídos onde a ordem de inicialização dos nós não é garantida.</li>
 * <li><b>Framing:</b> Implementação transparente do protocolo "Length-Prefix" (4 bytes de tamanho + payload),
 * resolvendo o problema de fragmentação de stream TCP.</li>
 * <li><b>Serialização:</b> Integração direta com a camada de serialização JSON.</li>
 * </ol>
 */
public class SocketConnection implements Closeable {

    // --- Network Resources ---

    /**
     * The underlying TCP socket used for network communication.
     */
    private final Socket socket;

    /**
     * The raw output stream for writing bytes to the network.
     * Wrapped by {@link DataOutputStream} during message sending.
     */
    private final OutputStream outputStream;

    /**
     * The raw input stream for reading bytes from the network.
     * Wrapped by {@link DataInputStream} during message reception.
     */
    private final InputStream inputStream;

    // --- Serialization ---

    /**
     * The serializer strategy used to convert objects to/from byte arrays (e.g., JSON).
     */
    private final MessageSerializer serializer;

    /** Número máximo de tentativas de ligação antes de desistir (Fail-fast). */
    private static final int MAX_RETRIES = 5; 
    
    /** Janela de espera (backoff) linear entre tentativas (em milissegundos). */
    private static final long RETRY_DELAY_MS = 1000; 

    /**
     * Construtor para clientes (Active Open).
     * Tenta estabelecer uma conexão TCP com um servidor, aplicando lógica de retry.
     * <p>
     * Este comportamento é vital quando o processo Coordenador inicia antes das Interseções estarem
     * prontas para aceitar conexões ({@code accept()}).
     *
     * @param host Endereço do nó de destino (ex: "localhost").
     * @param port Porta de serviço.
     * @throws IOException Se a conexão falhar após todas as {@code MAX_RETRIES} tentativas.
     * @throws UnknownHostException Se o DNS não resolver o hostname.
     * @throws InterruptedException Se a thread for interrompida durante o sleep de retry.
     */
    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 (SYN -> SYN-ACK -> ACK)
                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) {
                // Common errors: "Connection refused" (server not up) or "Timeout" (firewall/network)
                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) {
                    // Blocking wait before next attempt
                    TimeUnit.MILLISECONDS.sleep(RETRY_DELAY_MS);
                }
            } catch (IOException e) {
                // Other IO errors
                 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 ---

        // Final validation
        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; // Propagate the root cause
            } else {
                 throw new IOException("Failed to connect after " + MAX_RETRIES + " attempts, reason unknown.");
            }
        }

        // Initialize streams
        this.socket = tempSocket;
        this.outputStream = socket.getOutputStream();
        this.inputStream = socket.getInputStream();
        this.serializer = SerializerFactory.createDefault();
    }

    /**
     * Construtor para servidores (Passive Open).
     * Envolve um socket já conectado (retornado por {@code serverSocket.accept()}).
     * Não necessita de retry logic pois a conexão física já existe.
     *
     * @param acceptedSocket O socket ativo retornado pelo SO.
     * @throws IOException Se falhar a obtenção dos streams de I/O.
     */
    public SocketConnection(Socket acceptedSocket) throws IOException {
        this.socket = acceptedSocket;
        this.outputStream = socket.getOutputStream();
        this.inputStream = socket.getInputStream();
        this.serializer = SerializerFactory.createDefault();
    }

    /**
     * Serializa e transmite uma mensagem através do canal.
     * <p>
     * Utiliza sincronização ({@code synchronized}) para garantir que escritas concorrentes
     * na mesma conexão não corrompem a stream de bytes (thread-safety).
     *
     * @param message O objeto de protocolo a enviar.
     * @throws IOException Se o socket estiver fechado ou ocorrer erro de escrita.
     */
    public synchronized 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 (Framing)
            DataOutputStream dataOut = new DataOutputStream(outputStream);
            dataOut.writeInt(data.length);
            dataOut.write(data);
            dataOut.flush();
            LockSupport.parkNanos(50000); // 50us
            
        } catch (SerializationException e) {
            throw new IOException("Failed to serialize message", e);
        }
    }

    /**
     * Bloqueia à espera de uma mensagem completa do socket.
     * <p>
     * Lê primeiro o cabeçalho de tamanho (4 bytes) e depois o payload exato,
     * garantindo que processa mensagens completas mesmo se chegarem fragmentadas em múltiplos pacotes TCP.
     *
     * @return O objeto {@link MessageProtocol} reconstruído.
     * @throws IOException Se a conexão for perdida (EOF) ou o stream corrompido.
     * @throws ClassNotFoundException Se o tipo desserializado não for encontrado no classpath.
     */
    public MessageProtocol receiveMessage() throws IOException, ClassNotFoundException {
       if (socket == null || !socket.isConnected()) {
            throw new IOException("Socket is not connected");
        }

        try {
           
            DataInputStream dataIn = new DataInputStream(inputStream);
            int length = dataIn.readInt();
            
            // Sanity check para evitar OutOfMemory em caso de corrupção de stream
            if (length <= 0 || length > 10_000_000) { // Max 10MB payload
                throw new IOException("Invalid message length: " + length);
            }
            
            // Ler dados exatos da mensagem
            byte[] data = new byte[length];
            dataIn.readFully(data);
            LockSupport.parkNanos(50000); // 50us
            
            // Deserialize do JSON - força o tipo concreto Message
            return serializer.deserialize(data, sd.model.Message.class);
            
        } catch (SerializationException e) {
            throw new IOException("Failed to deserialize message", e);
        }
    }

    /**
     * Encerra a conexão e liberta os descritores de ficheiro.
     * Operação idempotente.
     */
    @Override
    public void close() throws IOException {
        if (inputStream != null) inputStream.close();
        if (outputStream != null) outputStream.close();
        if (socket != null) socket.close();
    }

    /**
     * Verifica o estado operacional da conexão.
     * @return true se o socket está aberto e conectado.
     */
    public boolean isConnected() {
        return socket != null && socket.isConnected() && !socket.isClosed();
    }
}