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.Socket;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.time.Duration;
import java.util.Arrays;
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.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.logging.Level;
import java.util.logging.Logger;

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

    private final AgentConfig config;
    private final BiConsumer<String, String> learnedTlsFingerprintConsumer;
    private final ExecutorService executor = Executors.newThreadPerTaskExecutor(
            new NamedThreadFactory("proxylink-backend", true)
    );
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final Map<String, String> learnedTlsFingerprints = new ConcurrentHashMap<>();
    private final Map<String, BackendTunnel> activeTunnels = new ConcurrentHashMap<>();
    private final AtomicReference<String> lastEvent = new AtomicReference<>("not started");

    public BackendAgent(AgentConfig config) {
        this(config, (ignoredEndpoint, ignoredFingerprint) -> {
        });
    }

    public BackendAgent(AgentConfig config, BiConsumer<String, String> learnedTlsFingerprintConsumer) {
        this.config = config;
        this.learnedTlsFingerprintConsumer = learnedTlsFingerprintConsumer;
    }

    @Override
    public void start() {
        if (!running.compareAndSet(false, true)) {
            return;
        }
        for (AgentConfig.FrontendEndpointConfig frontend : config.frontends()) {
            executor.execute(() -> runReconnectLoop(frontend));
        }
        lastEvent.set("connecting to " + config.frontends().size() + " frontend endpoint(s)");
        LOGGER.info(() -> "Backend agent connecting to " + config.frontends().size() + " frontend endpoint(s)");
    }

    @Override
    public TunnelStatus status() {
        int connectedTunnels = 0;
        int activeStreams = 0;
        for (BackendTunnel 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;
        }
        for (BackendTunnel tunnel : activeTunnels.values()) {
            tunnel.close("backend shutting down");
        }
        activeTunnels.clear();
        executor.shutdownNow();
        lastEvent.set("backend stopped");
        LOGGER.info("Backend agent stopped");
    }

    private void runReconnectLoop(AgentConfig.FrontendEndpointConfig frontend) {
        long delayMillis = config.tunnel().reconnectInitialMillis();
        while (running.get()) {
            try {
                runSingleConnection(frontend);
                delayMillis = config.tunnel().reconnectInitialMillis();
            } catch (Exception e) {
                if (running.get()) {
                    lastEvent.set("connection to " + frontend.name() + " failed: " + e.getMessage());
                    LOGGER.log(Level.WARNING, "Backend tunnel connection to " + frontend.name() + " failed", e);
                }
            }

            if (running.get()) {
                sleepWithBackoff(frontend, delayMillis);
                delayMillis = Math.min(delayMillis * 2, config.tunnel().reconnectMaxMillis());
            }
        }
    }

    private void runSingleConnection(AgentConfig.FrontendEndpointConfig frontend)
            throws IOException, GeneralSecurityException {
        Socket socket = null;
        boolean handedToTunnel = false;
        try {
            AgentConfig.FrontendEndpointConfig pinnedFrontend = frontendWithLearnedPin(frontend);
            socket = TlsSocketFactory.createClientSocket(
                    config,
                    pinnedFrontend,
                    fingerprint -> learnTlsFingerprint(frontend, fingerprint)
            );
            FrameCodec codec = new FrameCodec(socket.getInputStream(), socket.getOutputStream());
            authenticateToFrontend(codec);

            BackendTunnel tunnel = new BackendTunnel(frontend, socket, codec);
            BackendTunnel previous = activeTunnels.put(frontend.name(), tunnel);
            if (previous != null) {
                previous.close("replaced by a new authenticated tunnel to " + frontend.name());
            }

            handedToTunnel = true;
            lastEvent.set("authenticated to frontend " + frontend.name());
            LOGGER.info(() -> "Backend tunnel authenticated to " + frontend.name()
                    + " at " + frontend.connectHost() + ":" + frontend.connectPort());
            tunnel.runReadLoop();
        } finally {
            if (!handedToTunnel && socket != null) {
                try {
                    socket.close();
                } catch (IOException ignored) {
                }
            }
        }
    }

    private AgentConfig.FrontendEndpointConfig frontendWithLearnedPin(AgentConfig.FrontendEndpointConfig frontend) {
        if (!frontend.pinnedCertificateSha256().isBlank()) {
            return frontend;
        }
        String pin = learnedTlsFingerprints.getOrDefault(frontend.name(), "");
        if (pin.isBlank()) {
            return frontend;
        }
        return new AgentConfig.FrontendEndpointConfig(
                frontend.name(),
                frontend.connectHost(),
                frontend.connectPort(),
                pin
        );
    }

    private void learnTlsFingerprint(AgentConfig.FrontendEndpointConfig frontend, String fingerprint) {
        if (fingerprint == null || fingerprint.isBlank()) {
            return;
        }
        if (learnedTlsFingerprints.putIfAbsent(frontend.name(), fingerprint) == null) {
            learnedTlsFingerprintConsumer.accept(frontend.name(), fingerprint);
            lastEvent.set("pinned frontend TLS certificate for " + frontend.name());
            LOGGER.info("Pinned frontend TLS certificate fingerprint for " + frontend.name() + ": " + fingerprint);
        }
    }

    private void authenticateToFrontend(FrameCodec codec) throws IOException {
        AuthPayloads.Hello hello = new AuthPayloads.Hello(
                ProtocolConstants.PROTOCOL_VERSION,
                "backend",
                config.nodeName(),
                AuthPayloads.randomNonce()
        );
        codec.writeFrame(new Frame(FrameType.HELLO, 0, AuthPayloads.encodeHello(hello)));

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

        AuthPayloads.Challenge challenge = AuthPayloads.decodeChallenge(challengeFrame.payload());
        byte[] response = AuthPayloads.computeResponse(config.tunnel().authToken(), hello, challenge);
        codec.writeFrame(new Frame(FrameType.AUTH_RESPONSE, 0, AuthPayloads.encodeResponse(new AuthPayloads.Response(response))));

        Frame result = codec.readFrame();
        if (result == null) {
            throw new ProtocolException("Authentication ended before AUTH_OK");
        }
        if (result.type() == FrameType.AUTH_FAILED) {
            throw new ProtocolException("Frontend rejected authentication: "
                    + new String(result.payload(), StandardCharsets.UTF_8));
        }
        if (result.type() != FrameType.AUTH_OK) {
            throw new ProtocolException("Expected AUTH_OK frame, got " + result.type());
        }
    }

    private void sleepWithBackoff(AgentConfig.FrontendEndpointConfig frontend, long delayMillis) {
        long jitter = ThreadLocalRandom.current().nextLong(0, Math.max(1, delayMillis / 4));
        long sleepMillis = delayMillis + jitter;
        LOGGER.info(() -> "Reconnecting backend tunnel to " + frontend.name() + " in " + sleepMillis + "ms");
        try {
            Thread.sleep(sleepMillis);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private final class BackendTunnel {
        private final AgentConfig.FrontendEndpointConfig frontend;
        private final Socket socket;
        private final FrameCodec codec;
        private final AtomicBoolean open = new AtomicBoolean(true);
        private final Map<Long, LocalStream> streams = new ConcurrentHashMap<>();
        private final ScheduledExecutorService heartbeat = Executors.newSingleThreadScheduledExecutor(
                new NamedThreadFactory("proxylink-backend-heartbeat")
        );
        private volatile long lastReadMillis = System.currentTimeMillis();

        private BackendTunnel(AgentConfig.FrontendEndpointConfig frontend, Socket socket, FrameCodec codec) {
            this.frontend = frontend;
            this.socket = socket;
            this.codec = codec;
        }

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

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

        void runReadLoop() {
            scheduleHeartbeat();
            try {
                while (open.get()) {
                    Frame frame = codec.readFrame();
                    if (frame == null) {
                        throw new IOException("Frontend tunnel reached EOF");
                    }
                    lastReadMillis = System.currentTimeMillis();
                    handleFrame(frame);
                }
            } catch (Exception e) {
                if (open.get()) {
                    LOGGER.log(Level.WARNING, "Backend tunnel read loop stopped for frontend=" + frontend.name(), e);
                }
            } finally {
                close("frontend tunnel disconnected");
                activeTunnels.remove(frontend.name(), 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();
            try {
                socket.close();
            } catch (IOException ignored) {
            }
            for (LocalStream stream : streams.values()) {
                stream.markClosed();
                stream.closeSocket();
            }
            streams.clear();
            lastEvent.set("closed backend tunnel to " + frontend.name() + ": " + reason);
            LOGGER.info(() -> "Closed backend tunnel to " + frontend.name() + " 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 for " + frontend.name() + "; 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_OPEN -> openBackendStream(frame);
                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 frontend: " + 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 frontend: " + frame.type());
            }
        }

        private void openBackendStream(Frame frame) throws IOException {
            StreamOpenPayload payload = StreamOpenPayload.decode(frame.payload());
            AgentConfig.RouteConfig route;
            try {
                route = config.requireRoute(payload.routeName());
            } catch (ProtocolException e) {
                sendReset(frame.streamId(), e.getMessage());
                return;
            }

            Socket backendSocket = new Socket();
            try {
                backendSocket.setTcpNoDelay(true);
                backendSocket.setKeepAlive(true);
                backendSocket.connect(
                        new InetSocketAddress(route.backendTargetHost(), route.backendTargetPort()),
                        config.tunnel().connectTimeoutMillis()
                );
                LocalStream stream = new LocalStream(frame.streamId(), backendSocket);
                streams.put(frame.streamId(), stream);
                executor.execute(() -> pumpLocalToTunnel(stream, "backend-target-" + route.name()));
                LOGGER.fine(() -> "Opened backend stream " + frame.streamId() + " route=" + route.name());
            } catch (IOException e) {
                try {
                    backendSocket.close();
                } catch (IOException ignored) {
                }
                sendReset(frame.streamId(), "backend target connection failed");
                LOGGER.log(Level.WARNING, "Unable to connect route " + route.name()
                        + " target " + route.backendTargetHost() + ":" + route.backendTargetPort(), e);
            }
        }

        private void writeToLocalStream(Frame frame) throws IOException {
            LocalStream stream = streams.get(frame.streamId());
            if (stream == null || stream.isClosed()) {
                sendReset(frame.streamId(), "unknown backend 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)));
            }
        }
    }
}
