package me.proxylink.plugin.tunnel;

import me.proxylink.common.AgentConfig;
import me.proxylink.common.AuthPayloads;
import me.proxylink.common.Frame;
import me.proxylink.common.FrameCodec;
import me.proxylink.common.FrameType;
import me.proxylink.common.ProtocolConstants;
import me.proxylink.common.ProtocolException;
import me.proxylink.common.StreamOpenPayload;

import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level;
import java.util.logging.Logger;

public final class FrontendAgent implements ManagedAgent {
    private static final Logger LOGGER = Logger.getLogger(FrontendAgent.class.getName());

    private final AgentConfig config;
    private final ExecutorService executor = Executors.newThreadPerTaskExecutor(
            new NamedThreadFactory("proxylink-frontend", true)
    );
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final Map<String, FrontendTunnel> activeTunnels = new ConcurrentHashMap<>();
    private final AtomicReference<String> lastEvent = new AtomicReference<>("not started");
    private final List<ServerSocket> routeServerSockets = new ArrayList<>();
    private volatile ServerSocket tunnelServerSocket;

    public FrontendAgent(AgentConfig config) {
        this.config = config;
    }

    @Override
    public void start() throws IOException, GeneralSecurityException {
        if (!running.compareAndSet(false, true)) {
            return;
        }

        tunnelServerSocket = TlsSocketFactory.createServerSocket(config);
        executor.execute(this::acceptBackendTunnels);

        for (AgentConfig.RouteConfig route : config.routes()) {
            ServerSocket routeSocket = new ServerSocket();
            routeSocket.setReuseAddress(true);
            routeSocket.bind(new InetSocketAddress(route.frontendBindHost(), route.frontendBindPort()));
            routeServerSockets.add(routeSocket);
            executor.execute(() -> acceptPlayers(route, routeSocket));
            LOGGER.info(() -> "Route " + route.name() + " listening for proxy connections on "
                    + route.frontendBindHost() + ":" + route.frontendBindPort());
        }

        LOGGER.info(() -> "Frontend tunnel listener running on "
                + config.tunnel().listenHost() + ":" + config.tunnel().listenPort());
        lastEvent.set("listening on " + config.tunnel().listenHost() + ":" + config.tunnel().listenPort());
    }

    @Override
    public TunnelStatus status() {
        int connectedTunnels = 0;
        int activeStreams = 0;
        for (FrontendTunnel tunnel : activeTunnels.values()) {
            if (tunnel.isOpen()) {
                connectedTunnels++;
                activeStreams += tunnel.streamCount();
            }
        }
        boolean connected = running.get() && connectedTunnels > 0;
        return new TunnelStatus(
                config.role(),
                config.nodeName(),
                running.get(),
                connected,
                connectedTunnels,
                activeStreams,
                config.tls().enabled(),
                lastEvent.get()
        );
    }

    @Override
    public void close() {
        if (!running.compareAndSet(true, false)) {
            return;
        }

        closeQuietly(tunnelServerSocket);
        for (ServerSocket serverSocket : routeServerSockets) {
            closeQuietly(serverSocket);
        }

        for (FrontendTunnel tunnel : activeTunnels.values()) {
            tunnel.close("frontend shutting down");
        }
        activeTunnels.clear();

        executor.shutdownNow();
        lastEvent.set("frontend stopped");
        LOGGER.info("Frontend agent stopped");
    }

    private void acceptBackendTunnels() {
        while (running.get()) {
            try {
                Socket socket = tunnelServerSocket.accept();
                socket.setTcpNoDelay(true);
                socket.setKeepAlive(true);
                executor.execute(() -> handleBackendTunnel(socket));
            } catch (IOException e) {
                if (running.get()) {
                    LOGGER.log(Level.WARNING, "Failed to accept backend tunnel", e);
                }
            }
        }
    }

    private void handleBackendTunnel(Socket socket) {
        try {
            FrameCodec codec = new FrameCodec(socket.getInputStream(), socket.getOutputStream());
            AuthPayloads.Hello hello = authenticateBackend(codec);
            FrontendTunnel tunnel = new FrontendTunnel(socket, codec, hello.nodeName());
            FrontendTunnel previous = activeTunnels.put(hello.nodeName(), tunnel);
            if (previous != null) {
                previous.close("replaced by a new authenticated backend tunnel");
            }

            LOGGER.info(() -> "Authenticated backend tunnel from node=" + hello.nodeName()
                    + " remote=" + socket.getRemoteSocketAddress());
            lastEvent.set("backend authenticated: " + hello.nodeName());
            tunnel.runReadLoop();
        } catch (Exception e) {
            lastEvent.set("backend tunnel setup failed: " + e.getMessage());
            LOGGER.log(Level.WARNING, "Backend tunnel closed during setup or read loop", e);
            closeQuietly(socket);
        }
    }

    private AuthPayloads.Hello authenticateBackend(FrameCodec codec) throws IOException {
        Frame helloFrame = codec.readFrame();
        if (helloFrame == null || helloFrame.type() != FrameType.HELLO || helloFrame.streamId() != 0) {
            throw new ProtocolException("Expected HELLO frame");
        }

        AuthPayloads.Hello hello = AuthPayloads.decodeHello(helloFrame.payload());
        if (hello.protocolVersion() != ProtocolConstants.PROTOCOL_VERSION) {
            throw new ProtocolException("Backend protocol version mismatch: " + hello.protocolVersion());
        }
        if (!"backend".equals(hello.role())) {
            throw new ProtocolException("Only backend agents may authenticate to the frontend listener");
        }

        AuthPayloads.Challenge challenge = new AuthPayloads.Challenge(AuthPayloads.randomNonce());
        codec.writeFrame(new Frame(FrameType.AUTH_CHALLENGE, 0, AuthPayloads.encodeChallenge(challenge)));

        Frame responseFrame = codec.readFrame();
        if (responseFrame == null || responseFrame.type() != FrameType.AUTH_RESPONSE || responseFrame.streamId() != 0) {
            throw new ProtocolException("Expected AUTH_RESPONSE frame");
        }

        AuthPayloads.Response response = AuthPayloads.decodeResponse(responseFrame.payload());
        byte[] expected = AuthPayloads.computeResponse(config.tunnel().authToken(), hello, challenge);
        if (!AuthPayloads.constantTimeEquals(expected, response.hmac())) {
            codec.writeFrame(new Frame(FrameType.AUTH_FAILED, 0, "authentication failed".getBytes(StandardCharsets.UTF_8)));
            throw new ProtocolException("Backend authentication failed for node=" + hello.nodeName());
        }

        codec.writeFrame(new Frame(FrameType.AUTH_OK, 0, "ok".getBytes(StandardCharsets.UTF_8)));
        return hello;
    }

    private void acceptPlayers(AgentConfig.RouteConfig route, ServerSocket serverSocket) {
        while (running.get()) {
            try {
                Socket playerSocket = serverSocket.accept();
                playerSocket.setTcpNoDelay(true);
                playerSocket.setKeepAlive(true);

                FrontendTunnel tunnel = selectTunnelForRoute(route);
                if (tunnel == null || !tunnel.isOpen()) {
                    LOGGER.warning(() -> "Rejecting player connection for route " + route.name()
                            + " because no matching backend tunnel is authenticated");
                    closeQuietly(playerSocket);
                    continue;
                }

                tunnel.openStream(route, playerSocket);
            } catch (IOException e) {
                if (running.get()) {
                    LOGGER.log(Level.WARNING, "Failed accepting player connection for route " + route.name(), e);
                }
            }
        }
    }

    private FrontendTunnel selectTunnelForRoute(AgentConfig.RouteConfig route) {
        if (!route.backendNode().isBlank()) {
            return activeTunnels.get(route.backendNode());
        }

        FrontendTunnel selected = null;
        int openTunnels = 0;
        for (FrontendTunnel tunnel : activeTunnels.values()) {
            if (!tunnel.isOpen()) {
                continue;
            }
            selected = tunnel;
            openTunnels++;
        }

        if (openTunnels == 1) {
            return selected;
        }
        if (openTunnels > 1) {
            lastEvent.set("route " + route.name() + " is ambiguous; set route." + route.name() + ".backendNode");
            LOGGER.warning(() -> "Route " + route.name()
                    + " has no backendNode but multiple backend tunnels are connected");
        }
        return null;
    }

    private void clearActiveTunnel(FrontendTunnel tunnel) {
        activeTunnels.remove(tunnel.backendNodeName, tunnel);
    }

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

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

    private final class FrontendTunnel {
        private final Socket socket;
        private final FrameCodec codec;
        private final String backendNodeName;
        private final AtomicBoolean open = new AtomicBoolean(true);
        private final AtomicLong nextStreamId = new AtomicLong(1);
        private final Map<Long, LocalStream> streams = new ConcurrentHashMap<>();
        private final ScheduledExecutorService heartbeat = Executors.newSingleThreadScheduledExecutor(
                new NamedThreadFactory("proxylink-frontend-heartbeat")
        );
        private volatile long lastReadMillis = System.currentTimeMillis();

        private FrontendTunnel(Socket socket, FrameCodec codec, String backendNodeName) {
            this.socket = socket;
            this.codec = codec;
            this.backendNodeName = backendNodeName;
        }

        boolean isOpen() {
            return open.get();
        }

        int streamCount() {
            return streams.size();
        }

        void openStream(AgentConfig.RouteConfig route, Socket playerSocket) {
            long streamId = nextStreamId.getAndAdd(2);
            try {
                LocalStream stream = new LocalStream(streamId, playerSocket);
                streams.put(streamId, stream);
                codec.writeFrame(new Frame(FrameType.STREAM_OPEN, streamId, new StreamOpenPayload(route.name()).encode()));
                executor.execute(() -> pumpLocalToTunnel(stream, "frontend-player-" + route.name()));
                LOGGER.fine(() -> "Opened stream " + streamId + " route=" + route.name());
            } catch (IOException e) {
                streams.remove(streamId);
                closeQuietly(playerSocket);
                LOGGER.log(Level.WARNING, "Failed to open stream for route " + route.name(), e);
            }
        }

        void runReadLoop() {
            scheduleHeartbeat();
            try {
                while (open.get()) {
                    Frame frame = codec.readFrame();
                    if (frame == null) {
                        throw new IOException("Backend tunnel reached EOF");
                    }
                    lastReadMillis = System.currentTimeMillis();
                    handleFrame(frame);
                }
            } catch (Exception e) {
                if (open.get()) {
                    LOGGER.log(Level.WARNING, "Frontend tunnel read loop stopped for backend=" + backendNodeName, e);
                }
            } finally {
                close("backend tunnel disconnected");
                clearActiveTunnel(this);
            }
        }

        void close(String reason) {
            if (!open.compareAndSet(true, false)) {
                return;
            }

            try {
                codec.writeFrame(new Frame(FrameType.GOAWAY, 0, reason.getBytes(StandardCharsets.UTF_8)));
            } catch (IOException ignored) {
            }

            heartbeat.shutdownNow();
            closeQuietly(socket);
            for (LocalStream stream : streams.values()) {
                stream.markClosed();
                stream.closeSocket();
            }
            streams.clear();
            lastEvent.set("closed backend tunnel: " + reason);
            LOGGER.info(() -> "Closed frontend tunnel for backend=" + backendNodeName + " reason=" + reason);
        }

        private void scheduleHeartbeat() {
            long interval = config.tunnel().heartbeatIntervalMillis();
            heartbeat.scheduleAtFixedRate(() -> {
                if (!open.get()) {
                    return;
                }
                long idleMillis = System.currentTimeMillis() - lastReadMillis;
                long disconnectAfterMillis = disconnectAfterMillis(interval);
                if (idleMillis > disconnectAfterMillis) {
                    close("heartbeat timeout after " + Duration.ofMillis(idleMillis));
                    return;
                }
                if (idleMillis > interval * 2) {
                    lastEvent.set("heartbeat delayed " + idleMillis + "ms; tunnel still open");
                }
                try {
                    codec.writeFrame(Frame.empty(FrameType.PING, 0));
                } catch (IOException e) {
                    close("failed to write heartbeat ping");
                }
            }, interval, interval, TimeUnit.MILLISECONDS);
        }

        private long disconnectAfterMillis(long interval) {
            long missedLimit = interval * config.tunnel().heartbeatMissesBeforeDisconnect();
            return Math.max(interval, Math.min(config.tunnel().heartbeatTimeoutMillis(), missedLimit));
        }

        private void handleFrame(Frame frame) throws IOException {
            switch (frame.type()) {
                case PING -> codec.writeFrame(Frame.empty(FrameType.PONG, 0));
                case PONG -> {
                }
                case STREAM_DATA -> writeToLocalStream(frame);
                case STREAM_CLOSE, STREAM_RESET -> closeStream(frame.streamId(), false, frame.type(), "remote closed");
                case ERROR -> LOGGER.warning(() -> "Tunnel error from backend: " + new String(frame.payload(), StandardCharsets.UTF_8));
                case GOAWAY -> close("remote goaway: " + new String(frame.payload(), StandardCharsets.UTF_8));
                default -> throw new ProtocolException("Unexpected frame from backend: " + frame.type());
            }
        }

        private void writeToLocalStream(Frame frame) throws IOException {
            LocalStream stream = streams.get(frame.streamId());
            if (stream == null || stream.isClosed()) {
                sendReset(frame.streamId(), "unknown frontend stream");
                return;
            }
            stream.write(frame.payload());
        }

        private void pumpLocalToTunnel(LocalStream stream, String label) {
            byte[] buffer = new byte[ProtocolConstants.STREAM_DATA_CHUNK_BYTES];
            try (InputStream input = stream.socket().getInputStream()) {
                int read;
                while (open.get() && !stream.isClosed() && (read = input.read(buffer)) != -1) {
                    byte[] payload = Arrays.copyOf(buffer, read);
                    codec.writeFrame(new Frame(FrameType.STREAM_DATA, stream.streamId(), payload));
                }
                closeStream(stream.streamId(), true, FrameType.STREAM_CLOSE, label + " eof");
            } catch (IOException e) {
                closeStream(stream.streamId(), true, FrameType.STREAM_RESET, label + " io error");
            }
        }

        private void closeStream(long streamId, boolean notifyRemote, FrameType closeType, String reason) {
            LocalStream stream = streams.remove(streamId);
            if (stream == null || !stream.markClosed()) {
                return;
            }
            stream.closeSocket();
            if (notifyRemote && open.get()) {
                try {
                    codec.writeFrame(new Frame(closeType, streamId, reason.getBytes(StandardCharsets.UTF_8)));
                } catch (IOException e) {
                    close("failed to send stream close");
                }
            }
        }

        private void sendReset(long streamId, String reason) throws IOException {
            if (open.get()) {
                codec.writeFrame(new Frame(FrameType.STREAM_RESET, streamId, reason.getBytes(StandardCharsets.UTF_8)));
            }
        }
    }
}
