WebClient/src/main/java/org/openautonomousconnection/webclient/ui/BrowserTab.java

package org.openautonomousconnection.webclient.ui;

import javafx.application.Platform;
import javafx.concurrent.Worker;
import javafx.embed.swing.JFXPanel;
import javafx.scene.Scene;
import javafx.scene.control.ContextMenu;
import javafx.scene.control.MenuItem;
import javafx.scene.input.Clipboard;
import javafx.scene.input.ClipboardContent;
import javafx.scene.input.ContextMenuEvent;
import javafx.scene.web.WebEngine;
import javafx.scene.web.WebHistory;
import javafx.scene.web.WebView;
import org.openautonomousconnection.luascript.security.LuaExecutionPolicy;
import org.openautonomousconnection.oacswing.component.OACPanel;
import org.openautonomousconnection.webclient.ClientImpl;
import org.openautonomousconnection.webclient.lua.WebLogger;
import org.openautonomousconnection.webclient.settings.FxEngine;
import org.w3c.dom.Document;

import javax.swing.*;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.awt.*;
import java.io.File;
import java.io.IOException;
import java.io.StringWriter;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.time.Instant;
import java.util.Base64;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

/**
 * Logical browser tab: stable UI key + embedded JavaFX WebView.
 *
 * <p>This class merges the previous {@code BrowserTab} + {@code TabView} into one component.</p>
 */
public final class BrowserTab extends OACPanel {

    private final String key;

    private final AtomicBoolean initialized = new AtomicBoolean(false);
    private final Consumer<String> onLocationChanged;
    private final WebLogger webLogger;
    private final ClientImpl protocolClient;

    private final boolean luaEnabled;
    private final LuaExecutionPolicy luaPolicy;

    private volatile Runnable openInNewTab;

    private JFXPanel fxPanel;
    private WebView webView;
    private WebEngine engine;
    private volatile FxEngine luaEngine;

    /**
     * Creates a browser tab.
     *
     * @param key              stable UI key (must match CardLayout key and titlebar tab title)
     * @param initialUrl       initial URL (used for logger context)
     * @param onLocationChange callback invoked on URL changes
     * @param luaEnabled       whether Lua is enabled for this tab
     * @param luaPolicy        execution policy for Lua
     * @param protocolClient   protocol client used for OAC network requests
     */
    public BrowserTab(String key,
                      String initialUrl,
                      Consumer<String> onLocationChange,
                      boolean luaEnabled,
                      LuaExecutionPolicy luaPolicy,
                      ClientImpl protocolClient) {
        super();
        this.key = Objects.requireNonNull(key, "key");
        this.onLocationChanged = Objects.requireNonNull(onLocationChange, "onLocationChange");
        this.protocolClient = Objects.requireNonNull(protocolClient, "protocolClient");
        this.webLogger = new WebLogger(Objects.requireNonNull(initialUrl, "initialUrl"), protocolClient);
        this.luaEnabled = luaEnabled;
        this.luaPolicy = Objects.requireNonNull(luaPolicy, "luaPolicy");
        setLayout(new BorderLayout());
    }

    private static String serializeDom(Document doc) throws Exception {
        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer t = tf.newTransformer();
        t.setOutputProperty(OutputKeys.METHOD, "html");
        t.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "yes");
        t.setOutputProperty(OutputKeys.ENCODING, "UTF-8");
        t.setOutputProperty(OutputKeys.INDENT, "yes");

        StringWriter sw = new StringWriter(64 * 1024);
        t.transform(new DOMSource(doc), new StreamResult(sw));
        return sw.toString();
    }

    private static String normalizeMime(String contentType) {
        String ct = (contentType == null || contentType.isBlank()) ? "application/octet-stream" : contentType.trim();
        int semi = ct.indexOf(';');
        String base = (semi >= 0 ? ct.substring(0, semi) : ct).trim();
        return base.isEmpty() ? "application/octet-stream" : base;
    }

    private static boolean isHtml(String contentType) {
        String ct = normalizeMime(contentType).toLowerCase(Locale.ROOT);
        return ct.equals("text/html") || ct.equals("application/xhtml+xml");
    }

    private static boolean isText(String contentType) {
        String ct = normalizeMime(contentType).toLowerCase(Locale.ROOT);
        return ct.startsWith("text/") || ct.equals("application/json") || ct.equals("application/xml") || ct.endsWith("+json") || ct.endsWith("+xml");
    }

    private static boolean isImage(String contentType) {
        String ct = normalizeMime(contentType).toLowerCase(Locale.ROOT);
        return ct.startsWith("image/");
    }

    private static boolean isPdf(String contentType) {
        String ct = normalizeMime(contentType).toLowerCase(Locale.ROOT);
        return ct.equals("application/pdf");
    }

    private static Charset charsetFromContentType(String contentType, Charset def) {
        if (contentType == null) return def;
        String[] parts = contentType.split(";");
        for (String p : parts) {
            String s = p.trim();
            if (s.toLowerCase(Locale.ROOT).startsWith("charset=")) {
                String name = s.substring("charset=".length()).trim();
                try {
                    return Charset.forName(name);
                } catch (Exception ignored) {
                    return def;
                }
            }
        }
        return def;
    }

    private static String extractFilenameFromContentDisposition(Map<String, String> headers) {
        if (headers == null || headers.isEmpty()) return null;

        String cd = null;
        for (Map.Entry<String, String> e : headers.entrySet()) {
            if (e.getKey() != null && e.getKey().equalsIgnoreCase("content-disposition")) {
                cd = e.getValue();
                break;
            }
        }
        if (cd == null || cd.isBlank()) return null;

        String lower = cd.toLowerCase(Locale.ROOT);
        int fn = lower.indexOf("filename=");
        if (fn < 0) return null;

        String v = cd.substring(fn + "filename=".length()).trim();
        if (v.startsWith("\"")) {
            int end = v.indexOf('"', 1);
            if (end > 1) return v.substring(1, end);
            return null;
        }
        int semi = v.indexOf(';');
        if (semi >= 0) v = v.substring(0, semi).trim();
        return v.isBlank() ? null : v;
    }

    private static String sanitizeFilename(String name) {
        String s = name.replace('\\', '_').replace('/', '_');
        s = s.replace("..", "_");
        s = s.replace(':', '_').replace('*', '_').replace('?', '_').replace('"', '_')
                .replace('<', '_').replace('>', '_').replace('|', '_');
        return s.isBlank() ? "download.bin" : s;
    }

    private static String extensionFromContentType(String contentType) {
        String ct = normalizeMime(contentType).toLowerCase(Locale.ROOT);

        if (ct.equals("application/pdf")) return ".pdf";
        if (ct.equals("application/zip")) return ".zip";
        if (ct.equals("application/x-7z-compressed")) return ".7z";
        if (ct.equals("application/x-rar-compressed")) return ".rar";
        if (ct.equals("application/gzip")) return ".gz";
        if (ct.equals("application/json")) return ".json";
        if (ct.equals("application/xml") || ct.endsWith("+xml")) return ".xml";

        if (ct.startsWith("image/")) {
            int slash = ct.indexOf('/');
            if (slash > 0 && slash < ct.length() - 1) {
                String ext = ct.substring(slash + 1).trim();
                if (!ext.isEmpty()) return "." + ext;
            }
        }

        if (ct.startsWith("text/")) return ".txt";
        return ".bin";
    }

    /**
     * Returns the stable tab key.
     *
     * @return key
     */
    public String getKey() {
        return key;
    }

    /**
     * Returns the protocol client used by this tab.
     *
     * @return protocol client
     */
    public ClientImpl getProtocolClient() {
        return protocolClient;
    }

    /**
     * Sets callback for opening current page in a new tab.
     *
     * @param callback callback
     */
    public void setOpenInNewTabCallback(Runnable callback) {
        this.openInNewTab = callback;
    }

    @Override
    public void addNotify() {
        super.addNotify();

        if (!initialized.compareAndSet(false, true)) {
            return;
        }

        fxPanel = new JFXPanel();
        add(fxPanel, BorderLayout.CENTER);

        Platform.runLater(() -> {
            webView = new WebView();
            webView.setContextMenuEnabled(false);

            engine = webView.getEngine();
            engine.setJavaScriptEnabled(false);

            engine.locationProperty().addListener((obs, oldV, newV) -> fireLocationChanged(newV));
            engine.getLoadWorker().stateProperty().addListener((obs, oldS, newS) -> {
                if (newS == Worker.State.RUNNING || newS == Worker.State.SUCCEEDED) {
                    fireLocationChanged(engine.getLocation());
                }
            });

            installCustomContextMenu();

            if (luaEnabled) {
                luaEngine = new FxEngine(engine, webView, luaPolicy, webLogger);
                luaEngine.install();
            } else {
                luaEngine = null;
            }

            fxPanel.setScene(new Scene(webView));
        });
    }

    /**
     * Loads a URL.
     *
     * @param url URL
     */
    public void loadUrl(String url) {
        String target = Objects.requireNonNull(url, "url").trim();
        if (target.isEmpty()) return;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e != null) e.load(target);
        });
    }

    /**
     * Reloads the current page.
     */
    public void reload() {
        String current = getEngineLocation();
        if (current == null || current.isBlank()) return;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e != null) e.reload();
        });
    }

    /**
     * Goes back in history if possible.
     */
    public void goBack() {
        if (peekHistoryTarget(-1) == null) return;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e == null) return;
            WebHistory h = e.getHistory();
            if (h.getCurrentIndex() > 0) h.go(-1);
        });
    }

    /**
     * Goes forward in history if possible.
     */
    public void goForward() {
        if (peekHistoryTarget(+1) == null) return;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e == null) return;
            WebHistory h = e.getHistory();
            if (h.getCurrentIndex() < h.getEntries().size() - 1) h.go(1);
        });
    }

    /**
     * Returns current location if known.
     *
     * @return URL or null
     */
    public String getLocationUrl() {
        return getEngineLocation();
    }

    // -------------------- Stream render/save helpers --------------------

    /**
     * Returns current engine location.
     *
     * @return url or null
     */
    public String getEngineLocation() {
        WebEngine e = engine;
        return e != null ? e.getLocation() : null;
    }

    /**
     * Applies HTML live to the current tab (no saving).
     *
     * @param html html text
     */
    public void applyHtml(String html) {
        String content = html == null ? "" : html;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e != null) e.loadContent(content, "text/html");
        });
    }

    /**
     * Returns the current rendered HTML (DOM serialization).
     *
     * @return current html (never null)
     */
    public String getCurrentHtml() {
        CompletableFuture<String> fut = new CompletableFuture<>();
        Platform.runLater(() -> {
            try {
                WebEngine e = engine;
                if (e == null) {
                    fut.complete("");
                    return;
                }
                Document doc = e.getDocument();
                if (doc == null) {
                    fut.complete("");
                    return;
                }
                fut.complete(serializeDom(doc));
            } catch (Throwable t) {
                fut.completeExceptionally(t);
            }
        });

        try {
            return fut.join();
        } catch (Throwable t) {
            return "";
        }
    }

    /**
     * Receives a fully assembled stream payload from the protocol layer and renders or saves it.
     *
     * <p>Rules:
     * <ul>
     *   <li>Renderable: HTML/text/images/pdf are rendered directly (no wrapper pages for non-html).</li>
     *   <li>Not renderable: opens "Save As" dialog.</li>
     *   <li>If user cancels "Save As": content is shown raw in the browser via data: URL.</li>
     * </ul>
     *
     * @param requestId   request correlation id
     * @param tabId       protocol tab id
     * @param pageId      protocol page id
     * @param frameId     protocol frame id
     * @param statusCode  http-like status code
     * @param contentType mime type
     * @param headers     response headers
     * @param content     full payload bytes
     */
    public void handleStreamFinished(long requestId,
                                     long tabId,
                                     long pageId,
                                     long frameId,
                                     int statusCode,
                                     String contentType,
                                     Map<String, String> headers,
                                     byte[] content) {

        String ct = (contentType == null || contentType.isBlank())
                ? "application/octet-stream"
                : contentType.trim();

        byte[] data = (content == null) ? new byte[0] : content;

        // ---- Renderable types -> render without extra wrapper pages ----
        if (isHtml(ct)) {
            Charset cs = charsetFromContentType(ct, StandardCharsets.UTF_8);
            String html = new String(data, cs);
            Platform.runLater(() -> {
                WebEngine e = engine;
                if (e != null) e.loadContent(html, "text/html");
            });
            return;
        }

        if (isText(ct)) {
            Charset cs = charsetFromContentType(ct, StandardCharsets.UTF_8);
            String text = new String(data, cs);
            Platform.runLater(() -> {
                WebEngine e = engine;
                if (e != null) e.loadContent(text, "text/plain");
            });
            return;
        }

        if (isImage(ct) || isPdf(ct)) {
            renderRawDataUrl(ct, data);
            return;
        }

        // ---- Not renderable -> Save As; if cancelled -> raw data: URL ----
        String suggested = extractFilenameFromContentDisposition(headers);
        if (suggested == null || suggested.isBlank()) {
            String ext = extensionFromContentType(ct);
            suggested = "download_" + requestId + "_" + Instant.now().toEpochMilli() + ext;
        } else {
            suggested = sanitizeFilename(suggested);
        }

        showSaveAsDialogAndWriteBytes(suggested, ct, data);
    }

    /**
     * Releases resources.
     */
    public void dispose() {
        FxEngine le = luaEngine;
        luaEngine = null;
        if (le != null) {
            try {
                le.close();
            } catch (Exception ignored) {
                // Intentionally ignored
            }
        }

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e != null) {
                try {
                    e.load(null);
                } catch (Exception ignored) {
                    // Intentionally ignored
                }
            }
            engine = null;
            webView = null;
            if (fxPanel != null) {
                fxPanel.setScene(null);
            }
        });

        disconnectProtocolQuietly();
    }

    private void fireLocationChanged(String location) {
        if (location == null) return;
        String s = location.trim();
        if (s.isEmpty()) return;

        try {
            onLocationChanged.accept(s);
        } catch (Exception ignored) {
            // Must not break FX thread
        }
    }

    private void installCustomContextMenu() {
        final ContextMenu menu = new ContextMenu();

        MenuItem back = new MenuItem("Back");
        back.setOnAction(e -> SwingUtilities.invokeLater(this::goBack));

        MenuItem forward = new MenuItem("Forward");
        forward.setOnAction(e -> SwingUtilities.invokeLater(this::goForward));

        MenuItem reload = new MenuItem("Reload");
        reload.setOnAction(e -> SwingUtilities.invokeLater(this::reload));

        MenuItem copyLink = new MenuItem("Copy Link");
        copyLink.setOnAction(e -> {
            WebEngine e2 = engine;
            String loc = e2 != null ? e2.getLocation() : null;
            if (loc == null) return;
            ClipboardContent cc = new ClipboardContent();
            cc.putString(loc);
            Clipboard.getSystemClipboard().setContent(cc);
        });

        MenuItem openNewTab = new MenuItem("Open in New Tab");
        openNewTab.setOnAction(e -> {
            Runnable r = this.openInNewTab;
            if (r != null) SwingUtilities.invokeLater(r);
        });

        menu.getItems().addAll(back, forward, reload, copyLink, openNewTab);

        webView.addEventHandler(ContextMenuEvent.CONTEXT_MENU_REQUESTED, ev -> {
            menu.hide();
            menu.show(webView, ev.getScreenX(), ev.getScreenY());
            ev.consume();
        });
    }

    private String peekHistoryTarget(int delta) {
        WebEngine e = engine;
        if (e == null) return null;
        WebHistory h = e.getHistory();
        int idx = h.getCurrentIndex() + delta;
        if (idx < 0 || idx >= h.getEntries().size()) return null;
        return h.getEntries().get(idx).getUrl();
    }

    private void disconnectProtocolQuietly() {
        try {
            if (protocolClient.getClientServerConnection() != null) {
                protocolClient.getClientServerConnection().disconnect();
            }
        } catch (Exception ignored) {
            // Best-effort shutdown.
        }
        try {
            if (protocolClient.getClientINSConnection() != null) {
                protocolClient.getClientINSConnection().disconnect();
            }
        } catch (Exception ignored) {
            // Best-effort shutdown.
        }
    }

    private void showSaveAsDialogAndWriteBytes(String suggestedFilename, String contentType, byte[] data) {
        SwingUtilities.invokeLater(() -> {
            Window parent = SwingUtilities.getWindowAncestor(this);

            JFileChooser chooser = new JFileChooser();
            chooser.setDialogTitle("Save As");
            chooser.setSelectedFile(new File(suggestedFilename));

            int result = chooser.showSaveDialog(parent);
            if (result != JFileChooser.APPROVE_OPTION) {
                // Cancel -> show raw in browser
                renderRawDataUrl(contentType, data);
                return;
            }

            File file = chooser.getSelectedFile();
            if (file == null) {
                renderRawDataUrl(contentType, data);
                return;
            }

            if (file.isDirectory()) {
                JOptionPane.showMessageDialog(parent, "Please choose a file, not a directory.", "Save As", JOptionPane.WARNING_MESSAGE);
                renderRawDataUrl(contentType, data);
                return;
            }

            if (file.exists()) {
                int overwrite = JOptionPane.showConfirmDialog(
                        parent,
                        "File already exists. Overwrite?\n" + file.getAbsolutePath(),
                        "Confirm Overwrite",
                        JOptionPane.YES_NO_OPTION,
                        JOptionPane.WARNING_MESSAGE
                );
                if (overwrite != JOptionPane.YES_OPTION) {
                    renderRawDataUrl(contentType, data);
                    return;
                }
            }

            try {
                Files.write(file.toPath(), data);
            } catch (IOException ex) {
                JOptionPane.showMessageDialog(
                        parent,
                        "Failed to save file:\n" + ex.getMessage(),
                        "Save As",
                        JOptionPane.ERROR_MESSAGE
                );
                // On failure, still show raw so user can at least see bytes
                renderRawDataUrl(contentType, data);
            }
        });
    }

    private void renderRawDataUrl(String contentType, byte[] data) {
        String mime = normalizeMime(contentType);
        String b64 = Base64.getEncoder().encodeToString(data == null ? new byte[0] : data);
        String dataUrl = "data:" + mime + ";base64," + b64;

        Platform.runLater(() -> {
            WebEngine e = engine;
            if (e != null) e.load(dataUrl);
        });
    }

    public void bindProtocolClient() {
        protocolClient.getLibImpl().bindTab(this);
    }
}