WebClient/src/main/java/org/openautonomousconnection/webclient/ClientImpl.java

package org.openautonomousconnection.webclient;

import dev.unlegitdqrk.unlegitlibrary.event.Listener;
import lombok.Getter;
import org.openautonomousconnection.oacswing.component.OACOptionPane;
import org.openautonomousconnection.protocol.side.client.ProtocolClient;
import org.openautonomousconnection.protocol.side.client.events.ConnectedToProtocolINSServerEvent;
import org.openautonomousconnection.protocol.urlhandler.v1_0_1.beta.LibClientImpl_v1_0_1_B;
import org.openautonomousconnection.protocol.urlhandler.v1_0_1.beta.web.WebFlagInspector;
import org.openautonomousconnection.protocol.urlhandler.v1_0_1.beta.web.WebRequestContextProvider;
import org.openautonomousconnection.protocol.versions.v1_0_1.beta.WebPacketHeader;
import org.openautonomousconnection.webclient.ui.BrowserTab;

import java.awt.*;
import java.io.ByteArrayOutputStream;
import java.net.URL;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * WebClient Protocol implementation (v1.0.1).
 *
 * <p>Implements full stream assembly with strict correlation via:
 * requestId + tabId + pageId + frameId.</p>
 */
public final class ClientImpl extends ProtocolClient {

    private static final long MAX_STREAM_BYTES = 64L * 1024L * 1024L; // 64MB safety cap
    private static final int MAX_CONCURRENT_STREAMS = 256;

    @Getter
    private final LibImpl libImpl = new LibImpl();

    private final AtomicBoolean serverConnectionInitialized = new AtomicBoolean(false);
    private final Component dialogParent;
    private final Runnable onServerReady;

    public ClientImpl(Component dialogParent, Runnable onServerReady) {
        this.dialogParent = dialogParent;
        this.onServerReady = Objects.requireNonNull(onServerReady);
    }

    @Override
    public boolean trustINS(String caFingerprint) {
        Object[] options = {"Continue", "Cancel"};
        return OACOptionPane.showOptionDialog(
                dialogParent,
                "Fingerprint: " + caFingerprint + "\nContinue?",
                "INS Connection",
                OACOptionPane.YES_NO_OPTION,
                OACOptionPane.INFORMATION_MESSAGE,
                null,
                options,
                options[0]
        ) == 0;
    }

    @Override
    public boolean trustNewINSFingerprint(String oldCAFingerprint, String newCAFingerprint) {
        Object[] options = {"Continue", "Cancel"};
        return OACOptionPane.showOptionDialog(
                dialogParent,
                "Saved: " + oldCAFingerprint + "\nNew: " + newCAFingerprint + "\nContinue?",
                "INS Connection",
                OACOptionPane.YES_NO_OPTION,
                OACOptionPane.INFORMATION_MESSAGE,
                null,
                options,
                options[0]
        ) == 0;
    }

    @Listener
    public void onConnected(ConnectedToProtocolINSServerEvent event) {
        try {
            if (serverConnectionInitialized.compareAndSet(false, true)) {
                buildServerConnection(null, getProtocolBridge().getProtocolValues().ssl);
                onServerReady.run();
            }
        } catch (Exception e) {
            serverConnectionInitialized.set(false);
            throw new RuntimeException(e);
        }
    }

    private record StreamKey(long requestId, long tabId, long pageId, long frameId) {

        StreamKey(WebPacketHeader h) {
            this(h.getRequestId(), h.getTabId(), h.getPageId(), h.getFrameId());
        }
    }

    private static final class StreamState {

        private final int statusCode;
        private final String contentType;
        private final Map<String, String> headers;
        private final long declaredLength;

        private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        private int expectedSeq = 0;
        private long written = 0;
        private boolean ended = false;
        private boolean ok = true;

        StreamState(int statusCode,
                    String contentType,
                    Map<String, String> headers,
                    long declaredLength) {

            this.statusCode = statusCode;
            this.contentType = contentType == null ? "application/octet-stream" : contentType;
            this.headers = headers == null ? Map.of() : Map.copyOf(headers);
            this.declaredLength = declaredLength;
        }

        void append(int seq, byte[] data) {

            if (ended) throw new IllegalStateException("Chunk after end");
            if (seq != expectedSeq) throw new IllegalStateException("Out-of-order chunk");

            expectedSeq++;

            if (data == null || data.length == 0) return;

            written += data.length;
            if (written > MAX_STREAM_BYTES)
                throw new IllegalStateException("Stream exceeds limit");

            buffer.writeBytes(data);
        }

        void markEnd(boolean ok, String error) {
            this.ended = true;
            this.ok = ok;
        }

        byte[] finish() {
            if (!ok) return new byte[0];
            byte[] data = buffer.toByteArray();

            if (declaredLength > 0 && data.length != declaredLength) {
                // tolerated but can log if needed
            }

            return data;
        }
    }

    public final class LibImpl extends LibClientImpl_v1_0_1_B {

        private final WebRequestContextProvider provider = new WebRequestContextProvider.Default();
        private final WebFlagInspector inspector = new WebFlagInspector.Default();
        private final ConcurrentHashMap<StreamKey, StreamState> streams = new ConcurrentHashMap<>();

        private BrowserTab currentTab;

        @Override
        public void serverConnectionFailed(Exception exception) {
            getProtocolBridge().getLogger().exception("Failed to connect to server", exception);
            OACOptionPane.showMessageDialog(
                    dialogParent,
                    "Failed to connect to Server:\n" + exception.getMessage(),
                    "Server Connection",
                    OACOptionPane.ERROR_MESSAGE
            );
        }

        public void bindTab(BrowserTab tab) {
            this.currentTab = tab;
        }

        @Override
        public boolean isStream(WebPacketHeader header) {
            return inspector.isStream(header);
        }

        @Override
        public WebRequestContext contextFor(URL url) {
            return provider.contextFor(url);
        }

        @Override
        public void streamStart(WebPacketHeader header,
                                int statusCode,
                                String contentType,
                                Map<String, String> headers,
                                long totalLength) {

            if (streams.size() >= MAX_CONCURRENT_STREAMS) {
                throw new IllegalStateException("Too many concurrent streams");
            }

            StreamKey key = new StreamKey(header);
            streams.put(key, new StreamState(statusCode, contentType, headers, totalLength));
        }

        @Override
        public void streamChunk(WebPacketHeader header, int seq, byte[] data) {

            StreamState state = streams.get(new StreamKey(header));
            if (state == null) {
                throw new IllegalStateException("Chunk without streamStart");
            }

            state.append(seq, data);
        }

        @Override
        public void streamEnd(WebPacketHeader header, boolean ok, String error) {

            StreamState state = streams.get(new StreamKey(header));
            if (state != null) {
                state.markEnd(ok, error);
            }
        }

        @Override
        public void streamFinish(WebPacketHeader header, byte[] ignored) {

            StreamKey key = new StreamKey(header);
            StreamState state = streams.remove(key);
            if (state == null) return;

            byte[] content = state.finish();

            if (currentTab != null) {
                currentTab.handleStreamFinished(
                        header.getRequestId(),
                        header.getTabId(),
                        header.getPageId(),
                        header.getFrameId(),
                        state.statusCode,
                        state.contentType,
                        state.headers,
                        content
                );
            }
        }
    }
}