package jlama;

import com.formdev.flatlaf.FlatDarculaLaf;
import com.formdev.flatlaf.FlatIntelliJLaf;
import com.formdev.flatlaf.util.UIScale;
import org.jspecify.annotations.NonNull;

import javax.swing.*;
import javax.swing.border.EmptyBorder;
import javax.swing.tree.DefaultTreeCellRenderer;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreePath;
import java.awt.*;
import java.awt.event.WindowAdapter;
import java.awt.event.WindowEvent;
import java.io.IOException;
import java.net.URI;
import java.util.Vector;

public class Dashboard extends JFrame {
    public static final String WINDOW_TITLE = "LLM Demo";
    private static final int SMALL_GAP = UIScale.scale(5);
    private static final int MEDIUM_GAP = UIScale.scale(10);

    private static boolean lafInitialized;

    public static void initLaF() {
        if (!lafInitialized) {
            lafInitialized = true;
            try {
                if ("true".equals(System.getenv("JPROFILER_DARK_MODE"))) {
                    FlatDarculaLaf.setup();
                } else {
                    FlatIntelliJLaf.setup();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    public static void display() {
        initLaF();
        new Dashboard().setVisible(true);
    }

    private JLabel statusLabel;
    private JTree keywordTree;
    private JList<Article> articleList;
    private JLabel linkLabel;
    private JLabel titleLabel;
    private JTextPane summaryTextPane;

    private final KeywordContainer keywordContainer = new KeywordContainer(new KeywordContainer.NodeChangeListener() {
        @Override
        public void nodeInserted(KeywordNode node) {
            keywordTree.expandPath(new TreePath(node.getParent()));
        }

        @Override
        public void nodeUpdated(KeywordNode node) {
            if (new TreePath(node.getPath()).equals(keywordTree.getSelectionPath())) {
                updateDisplayedArticles();
            }
        }
    });
    private final LLMWorker llmWorker = new LLMWorker(keywordContainer, this::updateStatus);

    private Dashboard() {
        addComponents();
        setTitle(WINDOW_TITLE);
        setSize(new Dimension(UIScale.scale(600), UIScale.scale(500)));
        setLocation(UIScale.scale(100), UIScale.scale(100));

        setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE);
        addWindowListener(new WindowAdapter() {
            @Override
            public void windowClosing(WindowEvent e) {
                if (llmWorker.isAlive()) {
                    int result = JOptionPane.showConfirmDialog(Dashboard.this, "Do you want to stop the LLM demo?", WINDOW_TITLE, JOptionPane.YES_NO_OPTION, JOptionPane.QUESTION_MESSAGE);
                    if (result != JOptionPane.YES_OPTION) {
                        return;
                    }
                }
                System.exit(0);
            }

            @Override
            public void windowOpened(WindowEvent e) {
                llmWorker.start();
            }
        });
    }

    private void addComponents() {
        add(createHeaderPanel(), BorderLayout.NORTH);

        createArticleList();
        JSplitPane rightSplitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT, new JScrollPane(articleList), createArticlePanel());
        rightSplitPane.setDividerLocation(UIScale.scale(180));
        rightSplitPane.setResizeWeight(0.5);
        rightSplitPane.setContinuousLayout(true);

        createKeywordTree();
        JScrollPane keywordScrollPane = new JScrollPane(keywordTree);
        keywordScrollPane.setMinimumSize(new Dimension(UIScale.scale(150), UIScale.scale(50)));
        JSplitPane mainSplitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, keywordScrollPane, rightSplitPane);
        mainSplitPane.setDividerLocation(UIScale.scale(200));
        mainSplitPane.setResizeWeight(0.3);
        mainSplitPane.setContinuousLayout(true);

        add(mainSplitPane, BorderLayout.CENTER);

        updateStatus("Initializing ...");
    }

    private void updateStatus(String text) {
        statusLabel.setText(text);
    }

    private JPanel createHeaderPanel() {
        statusLabel = new JLabel();
        MultiLineTextDisplay textDisplay = new MultiLineTextDisplay(
                "This demo retrieves information about preprints from arxiv.org and " +
                "extracts keywords with a small local LLM. Articles are loaded in batches of " + LLMWorker.ARTICLE_BATCH_SIZE + ". " +
                "After each batch, similar keywords are consolidated by the LLM.\n\n" +
                "Check out the AI probe and the HTTP client probe in JProfiler for interesting data."
        );
        JPanel panel = new JPanel(new BorderLayout(0, MEDIUM_GAP));
        panel.setBorder(new EmptyBorder(MEDIUM_GAP, MEDIUM_GAP, MEDIUM_GAP, MEDIUM_GAP));
        panel.add(textDisplay, BorderLayout.NORTH);
        JPanel statusPanel = new JPanel(new BorderLayout(SMALL_GAP, 0));
        JLabel statusHeaderLabel = createBoldLabel();
        WaveProgressIcon.applyToLabel(statusHeaderLabel);
        statusHeaderLabel.setText("Status:");
        statusPanel.add(statusHeaderLabel, BorderLayout.WEST);
        statusPanel.add(statusLabel, BorderLayout.CENTER);
        panel.add(statusPanel, BorderLayout.SOUTH);
        return panel;
    }

    private void createKeywordTree() {
        DefaultTreeModel model = keywordContainer.getModel();
        keywordTree = new JTree(model);
        keywordTree.setRootVisible(false);
        keywordTree.setShowsRootHandles(false);
        keywordTree.expandPath(new TreePath(model.getRoot()));
        keywordTree.setCellRenderer(new DefaultTreeCellRenderer() {
            @Override
            public Component getTreeCellRendererComponent(JTree tree, Object value, boolean selected, boolean expanded, boolean leaf, int row, boolean hasFocus) {
                super.getTreeCellRendererComponent(tree, value, selected, expanded, leaf, row, hasFocus);
                if (value instanceof KeywordNode node) {
                    Keyword keyword = node.getKeyword();
                    if (keyword != null) {
                        String text = keyword.text();
                        if (keyword.resultCount() > -1) {
                            text += " [" + keyword.resultCount() + " hits]";
                        }
                        setText(text);
                    }
                }
                return this;
            }
        });
        keywordTree.addTreeSelectionListener(e -> updateDisplayedArticles());
    }

    private void updateDisplayedArticles() {
        KeywordNode node = (KeywordNode)keywordTree.getLastSelectedPathComponent();
        Vector<Article> data = new Vector<>();
        if (node != null && node.getKeyword() != null) {
            data.addAll(node.getArticles());
        }
        Article oldSelectedArticle = articleList.getSelectedValue();
        articleList.setListData(data);
        if (oldSelectedArticle != null) {
            articleList.setSelectedValue(oldSelectedArticle, true);
        }
    }

    private void createArticleList() {
        articleList = new JList<>();
        articleList.setSelectionMode(ListSelectionModel.SINGLE_SELECTION);
        articleList.setCellRenderer(new DefaultListCellRenderer() {
            @Override
            public Component getListCellRendererComponent(JList<?> list, Object value, int index, boolean isSelected, boolean cellHasFocus) {
                super.getListCellRendererComponent(list, value, index, isSelected, cellHasFocus);
                if (value instanceof Article article) {
                    setText(article.title());
                }
                return this;
            }
        });
        articleList.addListSelectionListener(e -> {
            Article article = articleList.getSelectedValue();
            if (article != null) {
                titleLabel.setText(article.title());
                linkLabel.setText(article.link().toString());
                summaryTextPane.setText(article.summary());
            } else {
                titleLabel.setText("");
                linkLabel.setText("");
                summaryTextPane.setText("");
            }
            summaryTextPane.setCaretPosition(0);
        });
    }

    private @NonNull JPanel createArticlePanel() {
        titleLabel = createBoldLabel();
        linkLabel = new JLabel();
        summaryTextPane = new JTextPane();
        summaryTextPane.setEditable(false);

        JButton openButton = new JButton("Open");
        openButton.addActionListener(e -> openLink());

        JPanel panel = new JPanel(new GridBagLayout());
        panel.setBorder(new EmptyBorder(MEDIUM_GAP, MEDIUM_GAP, SMALL_GAP, SMALL_GAP));
        GridBagConstraints gc = new GridBagConstraints();
        gc.gridx = 0;
        gc.gridy = 0;
        gc.anchor = GridBagConstraints.WEST;
        gc.insets = new Insets(0, 0, SMALL_GAP, SMALL_GAP);
        panel.add(new JLabel("Title:"), gc);

        gc.gridx = 1;
        gc.gridwidth = 2;
        gc.weightx = 1;
        gc.fill = GridBagConstraints.HORIZONTAL;
        panel.add(titleLabel, gc);

        gc.gridx = 0;
        gc.gridwidth = 1;
        gc.gridy = 1;
        gc.weightx = 0;
        gc.fill = GridBagConstraints.NONE;
        panel.add(new JLabel("Link:"), gc);

        gc.gridx = 1;
        gc.weightx = 1;
        gc.fill = GridBagConstraints.HORIZONTAL;
        panel.add(linkLabel, gc);

        gc.gridx = 2;
        gc.weightx = 0;
        gc.fill = GridBagConstraints.NONE;
        panel.add(openButton, gc);

        gc.gridx = 0;
        gc.gridy = 2;
        gc.weightx = 0;
        gc.weighty = 0;
        gc.fill = GridBagConstraints.NONE;
        gc.anchor = GridBagConstraints.NORTHWEST;
        panel.add(new JLabel("Summary:"), gc);

        gc.gridx = 1;
        gc.gridwidth = 2;
        gc.weightx = 1;
        gc.weighty = 1;
        gc.fill = GridBagConstraints.BOTH;
        panel.add(new JScrollPane(summaryTextPane), gc);

        return panel;
    }

    private JLabel createBoldLabel() {
        JLabel label = new JLabel();
        label.setFont(label.getFont().deriveFont(Font.BOLD));
        return label;
    }

    private void openLink() {
        String link = linkLabel.getText();
        if (!link.isEmpty()) {
            try {
                Desktop.getDesktop().browse(URI.create(link));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

}
