package com.dragonssh.xhttpdemo.core.tunnel;

import com.dragonssh.xhttpdemo.core.logger.SkStatus;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * Stable loopback TCP relay used between tun2socks and Trilead's dynamic SOCKS
 * forwarder.
 *
 * <p>The Android VPN always points to this relay. During an SSH/XHTTP reconnect,
 * only the relay's backend port changes. Keeping the public SOCKS listener alive
 * prevents tun2socks from treating a transient SSH outage as a fatal local proxy
 * failure and tearing down the VPN process.</p>
 */
final class PersistentSocksRelay {

    private static final String LOOPBACK_HOST = "127.0.0.1";
    private static final int ACCEPT_BACKLOG = 128;
    private static final int MAX_CLIENTS = 256;
    private static final int BACKEND_CONNECT_TIMEOUT_MS = 1500;
    private static final int BACKEND_RETRY_DELAY_MS = 150;
    private static final long BACKEND_WAIT_TIMEOUT_MS = 90_000L;
    private static final int COPY_BUFFER_SIZE = 32 * 1024;

    private final int listenPort;
    private final Object backendLock = new Object();
    private final Semaphore clientSlots = new Semaphore(MAX_CLIENTS);
    private final Set<Socket> openSockets = Collections.synchronizedSet(new HashSet<Socket>());

    private volatile boolean running;
    private volatile int backendPort = -1;
    private ServerSocket serverSocket;
    private Thread acceptThread;

    PersistentSocksRelay(int listenPort) {
        this.listenPort = listenPort;
    }

    synchronized void start() throws IOException {
        if (running) {
            return;
        }

        ServerSocket server = new ServerSocket();
        server.setReuseAddress(true);
        server.bind(new InetSocketAddress(LOOPBACK_HOST, listenPort), ACCEPT_BACKLOG);

        serverSocket = server;
        running = true;
        acceptThread = new Thread(this::acceptLoop, "vpn-socks-relay-accept");
        acceptThread.setDaemon(true);
        acceptThread.start();
    }

    boolean isRunning() {
        return running;
    }

    int getListenPort() {
        return listenPort;
    }

    void setBackendPort(int port) {
        synchronized (backendLock) {
            backendPort = port > 0 ? port : -1;
            backendLock.notifyAll();
        }
    }

    void clearBackendPort() {
        setBackendPort(-1);
    }

    synchronized void stop() {
        if (!running) {
            return;
        }

        running = false;
        clearBackendPort();

        ServerSocket server = serverSocket;
        serverSocket = null;
        closeQuietly(server);

        synchronized (openSockets) {
            for (Socket socket : openSockets) {
                closeQuietly(socket);
            }
            openSockets.clear();
        }

        Thread thread = acceptThread;
        acceptThread = null;
        if (thread != null) {
            thread.interrupt();
        }
    }

    private void acceptLoop() {
        while (running) {
            Socket client = null;
            try {
                ServerSocket server = serverSocket;
                if (server == null) {
                    break;
                }
                client = server.accept();
                client.setTcpNoDelay(true);
                client.setKeepAlive(true);

                if (!clientSlots.tryAcquire(1, TimeUnit.SECONDS)) {
                    closeQuietly(client);
                    continue;
                }

                final Socket accepted = client;
                openSockets.add(accepted);
                Thread worker = new Thread(() -> handleClient(accepted), "vpn-socks-relay-client");
                worker.setDaemon(true);
                worker.start();
            } catch (SocketException e) {
                if (running) {
                    SkStatus.logDebug("SOCKS relay accept failed: " + e.getMessage());
                }
                break;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                closeQuietly(client);
                break;
            } catch (IOException e) {
                closeQuietly(client);
                if (running) {
                    SkStatus.logDebug("SOCKS relay accept failed: " + e.getMessage());
                }
            }
        }
    }

    private void handleClient(Socket client) {
        Socket backend = null;
        try {
            backend = connectBackend();
            if (backend == null) {
                return;
            }

            openSockets.add(backend);
            relayBidirectional(client, backend);
        } catch (IOException ignored) {
            // A backend channel closing during reconnect is expected.
        } finally {
            closeQuietly(client);
            closeQuietly(backend);
            openSockets.remove(client);
            if (backend != null) {
                openSockets.remove(backend);
            }
            clientSlots.release();
        }
    }

    private Socket connectBackend() throws IOException {
        final long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(BACKEND_WAIT_TIMEOUT_MS);
        IOException lastError = null;

        while (running && System.nanoTime() < deadline) {
            int port = awaitBackendPort(deadline);
            if (port <= 0) {
                continue;
            }

            Socket socket = new Socket();
            try {
                socket.setTcpNoDelay(true);
                socket.setKeepAlive(true);
                socket.connect(new InetSocketAddress(LOOPBACK_HOST, port),
                        BACKEND_CONNECT_TIMEOUT_MS);
                return socket;
            } catch (IOException e) {
                lastError = e;
                closeQuietly(socket);

                try {
                    Thread.sleep(BACKEND_RETRY_DELAY_MS);
                } catch (InterruptedException interrupted) {
                    Thread.currentThread().interrupt();
                    return null;
                }
            }
        }

        if (lastError != null) {
            throw lastError;
        }
        return null;
    }

    private int awaitBackendPort(long deadlineNanos) {
        synchronized (backendLock) {
            while (running && backendPort <= 0) {
                long remainingNanos = deadlineNanos - System.nanoTime();
                if (remainingNanos <= 0) {
                    return -1;
                }

                long waitMillis = Math.max(1L,
                        Math.min(1000L, TimeUnit.NANOSECONDS.toMillis(remainingNanos)));
                try {
                    backendLock.wait(waitMillis);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    return -1;
                }
            }
            return backendPort;
        }
    }

    private static void relayBidirectional(Socket left, Socket right) throws IOException {
        final IOException[] upstreamError = new IOException[1];
        Thread upstream = new Thread(() -> {
            try {
                copy(left, right);
            } catch (IOException e) {
                upstreamError[0] = e;
            } finally {
                closeQuietly(left);
                closeQuietly(right);
            }
        }, "vpn-socks-relay-up");
        upstream.setDaemon(true);
        upstream.start();

        try {
            copy(right, left);
        } finally {
            closeQuietly(left);
            closeQuietly(right);
            try {
                upstream.join(1000L);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

        if (upstreamError[0] != null) {
            throw upstreamError[0];
        }
    }

    private static void copy(Socket source, Socket destination) throws IOException {
        InputStream in = source.getInputStream();
        OutputStream out = destination.getOutputStream();
        byte[] buffer = new byte[COPY_BUFFER_SIZE];
        int read;
        while ((read = in.read(buffer)) >= 0) {
            if (read == 0) {
                continue;
            }
            out.write(buffer, 0, read);
        }
    }

    private static void closeQuietly(Socket socket) {
        if (socket == null) {
            return;
        }
        try {
            socket.close();
        } catch (IOException ignored) {
        }
    }

    private static void closeQuietly(ServerSocket socket) {
        if (socket == null) {
            return;
        }
        try {
            socket.close();
        } catch (IOException ignored) {
        }
    }
}
