import json
import math
import sys
from pathlib import Path

from PySide6.QtCore import Qt, QSize, QStringListModel, QEvent
from PySide6.QtGui import (
    QColor, QFont, QIcon, QPainter, QPalette, QPixmap,
    QSyntaxHighlighter, QTextCharFormat, QTextCursor,
)
from PySide6.QtWidgets import (
    QApplication, QCompleter, QDialog, QFileDialog, QFrame, QHBoxLayout,
    QLabel, QListWidget, QMainWindow, QPushButton, QScrollArea,
    QSizePolicy, QTextBrowser, QTextEdit, QVBoxLayout, QWidget,
    QWidgetAction, QDockWidget, QCheckBox, QTabWidget, QListWidgetItem,
)

from abacus_core import AbacusCore


STYLE = """
QMainWindow, QWidget { background:#0d1117; color:#e6edf3; font-family:"Segoe UI"; }
QMenuBar { background:#111820; color:#c9d1d9; border-bottom:1px solid #202a36; }
QMenuBar::item:selected, QMenu::item:selected { background:#1b2531; }
QMenu { background:#111820; color:#c9d1d9; border:1px solid #202a36; }
QFrame#cell { background:#111820; border:1px solid #202a36; border-radius:12px; }
QFrame#cellError { background:#21171b; border:1px solid #7f3039; border-radius:12px; }
QLabel#prompt { color:#58a6ff; font-family:"JetBrains Mono"; font-weight:700; }
QLabel#resultTitle { color:#718096; font-size:11px; font-weight:700; }
QLabel#resultValue { color:#6ee7b7; font-family:"JetBrains Mono"; font-size:14px; font-weight:700; }
QLabel#errorValue { color:#ff8a8a; font-family:"JetBrains Mono"; }
QTextEdit#cellEditor { background:#0b0f14; color:#e6edf3; border:1px solid #344454;
 border-radius:6px; padding:2px 8px; font-family:"JetBrains Mono"; font-size:13px; }
QTextEdit#cellEditor:focus { border:1px solid #3b82f6; }
QPushButton { background:#2563eb; color:white; border:none; border-radius:7px; padding:7px 12px; font-weight:700; }
QPushButton:hover { background:#3475ef; }
QPushButton#small { background:#1b2531; color:#d7e0ea; border:1px solid #3a4a5c; padding:3px 10px; min-height:24px; font-size:12px; }
QPushButton#small:hover { background:#263649; color:#ffffff; }
QListWidget { background:#0d131a; color:#e6edf3; border:1px solid #202a36; border-radius:8px; }
QListWidget::item { color:#e6edf3; padding:7px; border-radius:6px; }
QListWidget::item:hover { background:#17212c; color:#ffffff; }
QListWidget::item:selected { background:#1d4f85; color:#ffffff; }
QListWidget::item:selected:!active { background:#26384a; color:#ffffff; }
QListWidget::indicator {
 width:16px; height:16px;
 border:1px solid #9aa8b8;
 border-radius:3px;
 background:#263342;
}
QListWidget::indicator:hover {
 border:1px solid #58a6ff;
 background:#33465a;
}
QListWidget::indicator:checked {
 background:#2563eb;
 border:1px solid #7db7ff;
}
QListWidget::indicator:checked:hover {
 background:#3475ef;
}
QCheckBox { color:#e6edf3; spacing:8px; }
QCheckBox::indicator { width:16px; height:16px; border:1px solid #718096; border-radius:3px; background:#111820; }
QCheckBox::indicator:hover { border:1px solid #58a6ff; background:#17212c; }
QCheckBox::indicator:checked { background:#2563eb; border:1px solid #58a6ff; }
QCheckBox::indicator:checked:hover { background:#3475ef; }
QScrollArea { border:none; background:#0d1117; }
QTabWidget::pane { border:1px solid #202a36; background:#0d1117; }
QTabBar::tab { background:#111820; color:#7d8b99; border:1px solid #202a36; padding:8px 15px; }
QTabBar::tab:selected { background:#17212c; color:#f0f6fc; }
QStatusBar { background:#0b0f14; color:#7d8b99; border-top:1px solid #202a36; }
"""


class AbacusHighlighter(QSyntaxHighlighter):
    def __init__(self, document):
        super().__init__(document)
        self.number = QTextCharFormat()
        self.number.setForeground(QColor("#79c0ff"))
        self.operator = QTextCharFormat()
        self.operator.setForeground(QColor("#ff7b72"))
        self.comment = QTextCharFormat()
        self.comment.setForeground(QColor("#6e7681"))
        self.comment.setFontItalic(True)
        self.keyword = QTextCharFormat()
        self.keyword.setForeground(QColor("#ffa657"))

    def highlightBlock(self, text):
        comment_pos = text.find("#")
        code = text if comment_pos < 0 else text[:comment_pos]
        if comment_pos >= 0:
            self.setFormat(comment_pos, len(text) - comment_pos, self.comment)

        i = 0
        while i < len(code):
            if code[i].isdigit() or (
                code[i] == "." and i + 1 < len(code) and code[i + 1].isdigit()
            ):
                start = i
                i += 1
                while i < len(code) and (code[i].isdigit() or code[i] == "."):
                    i += 1
                self.setFormat(start, i - start, self.number)
            else:
                i += 1

        for char in "+-*/=():[];":
            start = 0
            while True:
                pos = code.find(char, start)
                if pos < 0:
                    break
                self.setFormat(pos, 1, self.operator)
                start = pos + 1

        colon = code.find(":")
        if colon > 0:
            name = code[:colon].strip()
            if name:
                pos = code.find(name)
                self.setFormat(pos, len(name), self.keyword)


class CellEditor(QTextEdit):
    def __init__(self, run_callback):
        super().__init__()
        self.run_callback = run_callback
        self.completer = None

        self.setObjectName("cellEditor")
        self.setAcceptRichText(False)
        self.setLineWrapMode(QTextEdit.NoWrap)
        self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        self.setMinimumWidth(0)

        self._line_height = max(1, self.fontMetrics().lineSpacing())
        self._one_line_height = self._line_height + 8
        self.setFixedHeight(self._one_line_height)
        self.textChanged.connect(self._adjust_height)

    def _adjust_height(self):
        # Only explicit newline characters affect the editor height.
        # Width, word wrapping and QTextLayout are deliberately ignored.
        lines = max(1, self.toPlainText().count("\n") + 1)
        height = (
            self._one_line_height
            + (lines - 1) * self._line_height
        )
        if self.height() != height:
            self.setFixedHeight(height)
            self.updateGeometry()
            parent = self.parentWidget()
            if parent is not None:
                parent.updateGeometry()

    def set_completer(self, completer):
        self.completer = completer
        completer.setWidget(self)
        completer.setCompletionMode(QCompleter.PopupCompletion)
        completer.setCaseSensitivity(Qt.CaseInsensitive)

        popup = completer.popup()
        popup.setMinimumWidth(320)
        popup.setMaximumHeight(180)
        popup.setStyleSheet(
            "QListView {"
            "background:#111820;"
            "color:#e6edf3;"
            "border:1px solid #344454;"
            "padding:4px;"
            "}"
            "QListView::item {"
            "padding:5px 8px;"
            "}"
            "QListView::item:selected {"
            "background:#2563eb;"
            "color:white;"
            "}"
        )

        completer.activated.connect(self.insert_completion)

    def completion_prefix(self):
        cursor = self.textCursor()
        text = cursor.block().text()
        pos = cursor.positionInBlock()
        start = pos

        while start > 0 and (
            text[start - 1].isalnum() or text[start - 1] == "_"
        ):
            start -= 1

        return text[start:pos]

    def insert_completion(self, completion):
        prefix = self.completion_prefix()
        cursor = self.textCursor()

        if prefix:
            cursor.movePosition(
                QTextCursor.Left,
                QTextCursor.KeepAnchor,
                len(prefix),
            )

        cursor.insertText(completion)
        self.setTextCursor(cursor)

    def show_completion(self, all_names=False):
        if not self.completer:
            return

        prefix = self.completion_prefix()

        if not prefix and not all_names:
            self.completer.popup().hide()
            return

        self.completer.setCompletionPrefix(prefix)

        if self.completer.completionCount() == 0:
            self.completer.popup().hide()
            return

        rect = self.cursorRect()
        rect.setTop(rect.bottom() + 6)
        rect.setHeight(1)
        rect.setWidth(320)
        self.completer.complete(rect)

    def keyPressEvent(self, event):
        # Shift+Enter must work even while the completion popup is open.
        if event.key() in (Qt.Key_Return, Qt.Key_Enter) and (
            event.modifiers() & Qt.ShiftModifier
        ):
            if self.completer:
                self.completer.popup().hide()
            self.run_callback()
            return

        if self.completer and self.completer.popup().isVisible():
            if event.key() in (Qt.Key_Return, Qt.Key_Enter, Qt.Key_Tab):
                event.ignore()
                return

            if event.key() == Qt.Key_Escape:
                self.completer.popup().hide()
                return

        # Ctrl+Space: force completion list.
        if event.key() == Qt.Key_Space and (
            event.modifiers() & Qt.ControlModifier
        ):
            self.show_completion(True)
            return

        super().keyPressEvent(event)
        self.show_completion()


class NotebookCell(QFrame):
    def __init__(self, notebook, number=1, code=""):
        super().__init__()
        self.notebook = notebook
        self.number = number
        self.value = None
        self.error = None
        self.setObjectName("cell")
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)

        layout = QVBoxLayout(self)
        layout.setContentsMargins(8, 6, 8, 7)
        layout.setSpacing(4)

        header = QHBoxLayout()
        self.prompt = QLabel()
        self.prompt.setObjectName("prompt")
        header.addWidget(self.prompt)
        header.addStretch()

        for text, callback in (
            ("Run", self.run),
            ("+", self.add_after),
            ("×", self.delete),
        ):
            button = QPushButton(text)
            button.setObjectName("small")
            button.clicked.connect(callback)
            header.addWidget(button)

        layout.addLayout(header)

        self.editor = CellEditor(self.run)
        self.highlighter = AbacusHighlighter(self.editor.document())
        model = QStringListModel()
        self.completer = QCompleter(model, self)
        self.completion_model = model
        self.editor.set_completer(self.completer)
        self.editor.setPlainText(code)
        layout.addWidget(self.editor)

        self.result_title = QLabel("result:")
        self.result_title.setObjectName("resultTitle")
        layout.addWidget(self.result_title)

        self.result_label = QLabel("")
        self.result_label.setObjectName("resultValue")
        self.result_label.setTextInteractionFlags(Qt.TextSelectableByMouse)
        self.result_label.setSizePolicy(
            QSizePolicy.Expanding, QSizePolicy.Fixed
        )
        self.result_label.setFixedHeight(20)
        layout.addWidget(self.result_label)

        self.set_number(number)

    def set_number(self, number):
        self.number = number
        self.prompt.setText(f"In [{number}]")

    def set_completion_names(self, names):
        self.completion_model.setStringList(sorted(set(names)))

    def run(self):
        self.notebook.run_cell(self)

    def add_after(self):
        self.notebook.add_cell_after(self)

    def delete(self):
        self.notebook.delete_cell(self)

    def show_result(self, value):
        self.value = value
        self.error = None
        self.setObjectName("cell")
        self.style().unpolish(self)
        self.style().polish(self)
        self.result_title.setText("result:")
        self.result_label.setObjectName("resultValue")
        self.result_label.setText(str(value))

    def show_error(self, error):
        self.value = None
        self.error = str(error)
        self.setObjectName("cellError")
        self.style().unpolish(self)
        self.style().polish(self)
        self.result_title.setText("error:")
        self.result_label.setObjectName("errorValue")
        self.result_label.setText(str(error))

    def clear_output(self):
        self.value = None
        self.error = None
        self.result_title.setText("result:")
        self.result_label.setText("")

    def to_dict(self):
        return {
            "step": self.number,
            "code": self.editor.toPlainText(),
            "result": self.value,
            "error": self.error,
        }


class GraphWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.variables = {}

    def set_variables(self, variables):
        self.variables = dict(variables)
        self.update()

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.fillRect(self.rect(), QColor("#0b0f14"))

        area = self.rect().adjusted(60, 30, -25, -55)
        painter.setPen(QColor("#273442"))
        painter.drawRect(area)

        numeric = [
            (name, float(value))
            for name, value in self.variables.items()
            if isinstance(value, (int, float))
            and not isinstance(value, bool)
            and math.isfinite(float(value))
        ]

        if not numeric:
            painter.setPen(QColor("#647383"))
            painter.drawText(area, Qt.AlignCenter, "Select numeric variables")
            return

        lo = min(0.0, min(value for _, value in numeric))
        hi = max(0.0, max(value for _, value in numeric))
        if math.isclose(lo, hi):
            hi = lo + 1.0

        span = hi - lo
        zero_y = area.bottom() - ((0 - lo) / span) * area.height()

        painter.setPen(QColor("#344454"))
        painter.drawLine(
            area.left(), int(zero_y), area.right(), int(zero_y)
        )

        count = len(numeric)
        slot = area.width() / count
        bar_width = min(70, slot * 0.62)

        for index, (name, value) in enumerate(numeric):
            center_x = area.left() + slot * (index + 0.5)
            value_y = area.bottom() - ((value - lo) / span) * area.height()

            top = min(zero_y, value_y)
            height = max(2, abs(zero_y - value_y))

            painter.setBrush(QColor("#3b82f6"))
            painter.setPen(QColor("#58a6ff"))
            painter.drawRoundedRect(
                int(center_x - bar_width / 2),
                int(top),
                int(bar_width),
                int(height),
                5,
                5,
            )

            painter.setPen(QColor("#d6deea"))
            painter.drawText(
                int(center_x - slot / 2),
                int(top - 23 if value >= 0 else zero_y + height + 5),
                int(slot),
                18,
                Qt.AlignCenter,
                f"{value:g}",
            )

            painter.setPen(QColor("#8b98a7"))
            painter.drawText(
                int(center_x - slot / 2),
                area.bottom() + 10,
                int(slot),
                30,
                Qt.AlignCenter,
                name,
            )


class HelpDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Abacus - Help")
        self.resize(900, 760)

        layout = QVBoxLayout(self)

        browser = QTextBrowser()
        browser.setHtml(self._help_html())
        layout.addWidget(browser, 1)

        logo = QLabel()
        logo.setAlignment(Qt.AlignCenter)
        logo_path = Path(__file__).resolve().parent / "assets" / "abacus_logo.png"
        pixmap = QPixmap(str(logo_path))
        if not pixmap.isNull():
            logo.setPixmap(
                pixmap.scaled(
                    300, 300,
                    Qt.KeepAspectRatio,
                    Qt.SmoothTransformation,
                )
            )
        layout.addWidget(logo, 0, Qt.AlignCenter)

        close_button = QPushButton("Close")
        close_button.clicked.connect(self.accept)
        layout.addWidget(close_button, 0, Qt.AlignRight)

    @staticmethod
    def _help_html():
        return """
        <h2>How Abacus Works</h2>
        <p>
        Abacus is a notebook-style calculator. Calculations are written
        in cells. A cell may contain one or several lines. Lines inside
        a cell are executed sequentially from top to bottom and share
        the same variables and functions.
        </p>
        <h3>Variables</h3>
        <pre>anzahl_apfel = 10
gewicht_pro_apfel = 120
gesamt = anzahl_apfel * gewicht_pro_apfel</pre>
        <p>Leading and trailing spaces are ignored.</p>
        <h3>Expressions</h3>
        <p>Arithmetic operators are <b>+</b>, <b>-</b>, <b>*</b> and <b>/</b>.
        Parentheses can be used in expressions.</p>
        <h3>The result Variable</h3>
        <p>An ordinary expression stores its value in <b>result</b>.
        The cell displays the result of its last executed line.</p>
        <h3>Multiple Lines in One Cell</h3>
        <pre>a = 10
b = 20
c = a + b
c * 2</pre>
        <p>Lines are executed sequentially. Each line immediately sees
        variables created by previous lines.</p>
        <h3>Functions</h3>
        <pre>double: alpha * 2</pre>
        <p>Functions are shown in the Functions panel.</p>
        <h3>Comments</h3>
        <pre>alpha = 20 # working value</pre>
        <h3>Variables and Functions Panels</h3>
        <p>The Variables and Functions panels can be shown or hidden from
        View. Double-click an item to insert its name at the cursor.</p>
        <h3>Local and Global Items</h3>
        <p>An unchecked item is local. A checked item belongs to the global
        library and is loaded when Abacus starts.</p>
        <h3>Graph</h3>
        <p>The Graph tab displays selected numeric variables as bars.</p>
        <h3>Running a Cell</h3>
        <p>After execution, focus moves to the next cell. If there is no
        next cell, Abacus creates one automatically.</p>
        <h2>Keyboard Shortcuts</h2>
        <table width="100%" cellspacing="8">
          <tr><td><b>Shift + Enter</b></td><td>Run current cell and move to the next</td></tr>
          <tr><td><b>Ctrl + Space</b></td><td>Show autocomplete</td></tr>
          <tr><td><b>F5</b></td><td>Run all cells</td></tr>
          <tr><td><b>Ctrl + N</b></td><td>New notebook</td></tr>
          <tr><td><b>Ctrl + S</b></td><td>Export notebook</td></tr>
          <tr><td><b>Ctrl + O</b></td><td>Import notebook</td></tr>
          <tr><td><b>Ctrl + Q</b></td><td>Exit Abacus</td></tr>
          <tr><td><b>F1</b></td><td>Open Help</td></tr>
        </table>
        """


class AbacusNotebook(QMainWindow):
    def __init__(self):
        super().__init__()
        self.core = AbacusCore()
        self.cells = []
        self.current_cell = None
        self.graph_points = []
        self.global_library_path = (
            Path(__file__).resolve().with_name("abacus_globals.json")
        )
        self.global_variables = {}
        self.global_functions = {}
        self._updating_panels = False

        self._load_global_library()

        self.setWindowTitle("Abacus Notebook")
        icon_path = Path(__file__).resolve().parent / "assets" / "abacus_icon.ico"
        if icon_path.exists():
            self.setWindowIcon(QIcon(str(icon_path)))
        self.resize(1400, 850)
        self.setMinimumSize(1000, 650)

        self._build_menu()
        self._build_ui()
        self._build_sidebars()

        self.add_cell()
        self._refresh_sidebars()
        self._refresh_graph_variables()

    def _load_global_library(self):
        if not self.global_library_path.exists():
            return

        try:
            with self.global_library_path.open(
                "r", encoding="utf-8"
            ) as file:
                data = json.load(file)

            variables = data.get("variables", {})
            functions = data.get("functions", {})

            if isinstance(variables, dict):
                self.global_variables = variables

            if isinstance(functions, dict):
                self.global_functions = functions

            self.core.get_vars().update(self.global_variables)
            self.core.get_funcs().update(self.global_functions)

        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            self.global_variables = {}
            self.global_functions = {}

    def _save_global_library(self):
        data = {
            "format": "abacus-global-library",
            "version": 1,
            "variables": self.global_variables,
            "functions": self.global_functions,
        }

        try:
            with self.global_library_path.open(
                "w", encoding="utf-8"
            ) as file:
                json.dump(
                    data,
                    file,
                    ensure_ascii=False,
                    indent=2,
                )
        except (OSError, TypeError, ValueError):
            self.statusBar().showMessage(
                f"Cannot save global library: {self.global_library_path}"
            )

    def _set_global_item(self, kind, name, checked):
        if self._updating_panels:
            return

        if kind == "variable":
            source = self.core.get_vars()
            target = self.global_variables
        else:
            source = self.core.get_funcs()
            target = self.global_functions

        if name not in source:
            return

        if checked:
            target[name] = source[name]
        else:
            target.pop(name, None)

        self._save_global_library()

    def _insert_panel_item(self, item):
        if item is None:
            return

        name = item.data(Qt.UserRole)
        if not name:
            return

        cell = self.current_cell
        if cell is None or cell not in self.cells:
            if self.cells:
                cell = self.cells[0]
            else:
                self.add_cell()
                cell = self.cells[0]

        editor = cell.editor
        editor.setFocus()

        cursor = editor.textCursor()
        cursor.insertText(str(name))
        editor.setTextCursor(cursor)

    def eventFilter(self, watched, event):
        if event.type() == QEvent.FocusIn:
            for cell in self.cells:
                if watched is cell.editor:
                    self.current_cell = cell
                    break

        return super().eventFilter(watched, event)

    def _build_menu(self):
        menu = self.menuBar()
        file_menu = menu.addMenu("File")
        new_action = file_menu.addAction("New Notebook")
        new_action.setShortcut("Ctrl+N")
        new_action.triggered.connect(self.new_notebook)

        file_menu.addSeparator()

        export_action = file_menu.addAction("Export Notebook...")
        export_action.setShortcut("Ctrl+S")
        export_action.triggered.connect(self.export_notebook)

        import_action = file_menu.addAction("Import Notebook...")
        import_action.setShortcut("Ctrl+O")
        import_action.triggered.connect(self.import_notebook)

        file_menu.addSeparator()

        exit_action = file_menu.addAction("Exit")
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)

    def _build_view_menu(self):
        view_menu = self.menuBar().addMenu("View")

        variables_action = self.variables_dock.toggleViewAction()
        variables_action.setText("Variables")
        view_menu.addAction(variables_action)

        functions_action = self.functions_dock.toggleViewAction()
        functions_action.setText("Functions")
        view_menu.addAction(functions_action)

        view_menu.addSeparator()

        refresh_action = view_menu.addAction("Refresh panels")
        refresh_action.triggered.connect(self._refresh_all_panels)

        view_menu.addSeparator()
        library_action = view_menu.addAction("Global library file")
        library_action.triggered.connect(
            lambda: self.statusBar().showMessage(
                str(self.global_library_path)
            )
        )

        spacer = QWidget()
        spacer.setSizePolicy(
            QSizePolicy.Expanding,
            QSizePolicy.Preferred,
        )
        spacer_action = QWidgetAction(self.menuBar())
        spacer_action.setDefaultWidget(spacer)
        self.menuBar().addAction(spacer_action)

        help_menu = self.menuBar().addMenu("Help")
        help_action = help_menu.addAction("Abacus Help")
        help_action.setShortcut("F1")
        help_action.triggered.connect(self._show_help)

    def _show_help(self):
        HelpDialog(self).exec()

    def _build_ui(self):
        central = QWidget()
        root = QVBoxLayout(central)
        root.setContentsMargins(10, 10, 10, 10)

        toolbar = QHBoxLayout()
        title = QLabel("Abacus Notebook")
        title.setStyleSheet("font-size:18px;font-weight:700;color:#f0f6fc;")
        toolbar.addWidget(title)
        toolbar.addStretch()

        run_all = QPushButton("▶ Run All")
        run_all.setToolTip("Run all cells (F5)")
        run_all.clicked.connect(self.run_all)
        run_all.setShortcut("F5")
        toolbar.addWidget(run_all)

        add = QPushButton("+ Cell")
        add.clicked.connect(self._add_cell_from_button)
        toolbar.addWidget(add)
        root.addLayout(toolbar)

        self.tabs = QTabWidget()

        notebook_page = QWidget()
        notebook_layout = QVBoxLayout(notebook_page)
        notebook_layout.setContentsMargins(0, 0, 0, 0)

        self.scroll = QScrollArea()
        self.scroll.setWidgetResizable(True)

        self.notebook_widget = QWidget()
        self.notebook_layout = QVBoxLayout(self.notebook_widget)
        self.notebook_layout.setContentsMargins(8, 8, 8, 8)
        self.notebook_layout.setSpacing(7)
        self.notebook_layout.addStretch()

        self.scroll.setWidget(self.notebook_widget)
        notebook_layout.addWidget(self.scroll)
        self.tabs.addTab(notebook_page, "Notebook")

        graph_page = QWidget()
        graph_layout = QVBoxLayout(graph_page)
        graph_layout.setContentsMargins(10, 10, 10, 10)

        graph_header = QHBoxLayout()
        graph_header.addWidget(QLabel("Variables"))
        graph_header.addStretch()

        select_all = QPushButton("All")
        select_all.setObjectName("small")
        select_all.clicked.connect(lambda: self._set_all_graph_variables(True))
        graph_header.addWidget(select_all)

        clear_all = QPushButton("None")
        clear_all.setObjectName("small")
        clear_all.clicked.connect(lambda: self._set_all_graph_variables(False))
        graph_header.addWidget(clear_all)

        graph_layout.addLayout(graph_header)

        graph_body = QHBoxLayout()

        self.graph_variables = QWidget()
        self.graph_variables_layout = QVBoxLayout(self.graph_variables)
        self.graph_variables_layout.setContentsMargins(0, 0, 10, 0)
        self.graph_variables_layout.setSpacing(4)

        graph_scroll = QScrollArea()
        graph_scroll.setWidgetResizable(True)
        graph_scroll.setFixedWidth(240)
        graph_scroll.setWidget(self.graph_variables)
        graph_body.addWidget(graph_scroll)

        self.graph = GraphWidget()
        graph_body.addWidget(self.graph, 1)

        graph_layout.addLayout(graph_body, 1)
        self.tabs.addTab(graph_page, "Graph")

        root.addWidget(self.tabs, 1)

        self.setCentralWidget(central)
        self.statusBar().showMessage("Ready")

    def _refresh_graph_variables(self):
        old_checked = {
            checkbox.text(): checkbox.isChecked()
            for checkbox in self.graph_variables.findChildren(QCheckBox)
        }

        while self.graph_variables_layout.count():
            item = self.graph_variables_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        variables = dict(self.core.get_vars())

        for name, value in variables.items():
            if not isinstance(value, (int, float)) or isinstance(value, bool):
                continue

            checkbox = QCheckBox(name)
            checkbox.setChecked(old_checked.get(name, True))
            checkbox.toggled.connect(self._update_graph_from_selection)
            self.graph_variables_layout.addWidget(checkbox)

        self.graph_variables_layout.addStretch()
        self._update_graph_from_selection()

    def _update_graph_from_selection(self):
        selected = set()

        for checkbox in self.graph_variables.findChildren(QCheckBox):
            if checkbox.isChecked():
                selected.add(checkbox.text())

        variables = dict(self.core.get_vars())
        values = {
            name: value
            for name, value in variables.items()
            if name in selected
        }
        self.graph.set_variables(values)

    def _set_all_graph_variables(self, checked):
        for checkbox in self.graph_variables.findChildren(QCheckBox):
            checkbox.setChecked(checked)

    def _sync_global_library_from_core(self):
        variables = self.core.get_vars()
        functions = self.core.get_funcs()

        changed = False

        for name in list(self.global_variables):
            if name in variables and self.global_variables[name] != variables[name]:
                self.global_variables[name] = variables[name]
                changed = True
            elif name not in variables:
                self.global_variables.pop(name, None)
                changed = True

        for name in list(self.global_functions):
            if name in functions and self.global_functions[name] != functions[name]:
                self.global_functions[name] = functions[name]
                changed = True
            elif name not in functions:
                self.global_functions.pop(name, None)
                changed = True

        if changed:
            self._save_global_library()

    def _refresh_all_panels(self):
        self._refresh_sidebars()
        self._update_completion()
        self._refresh_graph_variables()
        self.statusBar().showMessage("Panels refreshed")

    def _make_panel_item(self, name, value, kind):
        item = QListWidgetItem()
        item.setData(Qt.UserRole, name)
        item.setData(Qt.UserRole + 1, kind)
        item.setText(f"{name} = {value}" if kind == "variable"
                     else f"{name} : {value}")
        item.setFlags(
            item.flags()
            | Qt.ItemIsUserCheckable
            | Qt.ItemIsSelectable
            | Qt.ItemIsEnabled
        )

        global_values = (
            self.global_variables
            if kind == "variable"
            else self.global_functions
        )
        item.setCheckState(
            Qt.Checked if name in global_values else Qt.Unchecked
        )
        return item

    def _refresh_sidebars(self):
        variables = dict(self.core.get_vars())
        functions = dict(self.core.get_funcs())

        self._updating_panels = True
        self.variables_list.blockSignals(True)
        self.functions_list.blockSignals(True)

        try:
            self.variables_list.clear()
            for name, value in variables.items():
                self.variables_list.addItem(
                    self._make_panel_item(name, value, "variable")
                )

            if not variables:
                empty = QListWidgetItem("(no variables)")
                empty.setForeground(QColor("#718096"))
                self.variables_list.addItem(empty)

            self.functions_list.clear()
            for name, value in functions.items():
                self.functions_list.addItem(
                    self._make_panel_item(name, value, "function")
                )

            if not functions:
                empty = QListWidgetItem("(no functions)")
                empty.setForeground(QColor("#718096"))
                self.functions_list.addItem(empty)

        finally:
            self.variables_list.blockSignals(False)
            self.functions_list.blockSignals(False)
            self._updating_panels = False


    def _add_cell_from_button(self, checked=False):
        self.add_cell()

    @staticmethod
    def _style_panel_list(widget):
        palette = widget.palette()
        palette.setColor(QPalette.Base, QColor("#0d131a"))
        palette.setColor(QPalette.Text, QColor("#e6edf3"))
        palette.setColor(QPalette.Button, QColor("#263342"))
        palette.setColor(QPalette.ButtonText, QColor("#e6edf3"))
        palette.setColor(QPalette.Highlight, QColor("#1d4f85"))
        palette.setColor(QPalette.HighlightedText, QColor("#ffffff"))
        widget.setPalette(palette)

    def _build_sidebars(self):
        self.variables_dock = QDockWidget("Variables", self)
        self.variables_list = QListWidget()
        self._style_panel_list(self.variables_list)
        self.variables_list.itemChanged.connect(
            self._variable_item_changed
        )
        self.variables_list.itemDoubleClicked.connect(
            self._insert_panel_item
        )
        self.variables_dock.setWidget(self.variables_list)
        self.variables_dock.setMinimumWidth(260)
        self.variables_dock.setAllowedAreas(
            Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea
        )
        self.addDockWidget(Qt.RightDockWidgetArea, self.variables_dock)

        self.functions_dock = QDockWidget("Functions", self)
        self.functions_list = QListWidget()
        self._style_panel_list(self.functions_list)
        self.functions_list.itemChanged.connect(
            self._function_item_changed
        )
        self.functions_list.itemDoubleClicked.connect(
            self._insert_panel_item
        )
        self.functions_dock.setWidget(self.functions_list)
        self.functions_dock.setMinimumWidth(260)
        self.functions_dock.setAllowedAreas(
            Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea
        )
        self.addDockWidget(Qt.RightDockWidgetArea, self.functions_dock)

        self.resizeDocks(
            [self.variables_dock, self.functions_dock],
            [300, 300],
            Qt.Vertical,
        )

        self.variables_dock.show()
        self.functions_dock.show()

        self._build_view_menu()

    def _variable_item_changed(self, item):
        if item is not None:
            self._set_global_item(
                "variable",
                item.data(Qt.UserRole),
                item.checkState() == Qt.Checked,
            )

    def _function_item_changed(self, item):
        if item is not None:
            self._set_global_item(
                "function",
                item.data(Qt.UserRole),
                item.checkState() == Qt.Checked,
            )


    def add_cell(self, after=None, code=""):
        cell = NotebookCell(self, code=code)

        index = len(self.cells) if after is None else self.cells.index(after) + 1
        self.cells.insert(index, cell)
        self.notebook_layout.insertWidget(index, cell)
        cell.editor.installEventFilter(self)
        self.current_cell = cell

        self._renumber()
        self._update_completion()
        cell.editor.setFocus()

    def add_cell_after(self, cell):
        self.add_cell(cell)

    def delete_cell(self, cell):
        if len(self.cells) == 1:
            cell.editor.clear()
            cell.clear_output()
            return

        index = self.cells.index(cell)
        self.cells.remove(cell)
        cell.setParent(None)
        cell.deleteLater()
        self._renumber()
        self._rebuild_graph()
        self._update_completion()
        self._refresh_graph_variables()

        if index < len(self.cells):
            self.cells[index].editor.setFocus()

    def _renumber(self):
        for index, cell in enumerate(self.cells, 1):
            cell.set_number(index)

    def _update_completion(self):
        names = list(self.core.get_vars()) + list(self.core.get_funcs())
        for cell in self.cells:
            cell.set_completion_names(names)

    def run_cell(self, cell, advance=True):
        code = cell.editor.toPlainText().strip()

        if not code:
            if advance:
                self._focus_next_cell(cell)
            return

        try:
            self.core.parse_input(code)

            # Always read the live dictionaries from the core after execution.
            # The GUI never keeps a second copy of variables/functions.
            variables = dict(self.core.get_vars())
            functions = dict(self.core.get_funcs())
            result_value = variables.get("result")

        except Exception as exc:
            cell.show_error(exc)
            self._refresh_sidebars()
            self._refresh_graph_variables()
            self.statusBar().showMessage(
                f"Error in In [{cell.number}]: {exc}"
            )
            return

        cell.show_result(result_value)

        # The Graph tab is based on the currently selected variables.
        self._refresh_sidebars()
        self._refresh_graph_variables()
        self._sync_global_library_from_core()

        self.statusBar().showMessage(
            f"In [{cell.number}]    result = {result_value}"
        )

        if advance:
            self._focus_next_cell(cell)

    def _focus_next_cell(self, cell):
        index = self.cells.index(cell)

        if index + 1 < len(self.cells):
            next_cell = self.cells[index + 1]
        else:
            self.add_cell()
            next_cell = self.cells[-1]

        next_cell.editor.setFocus()
        next_cell.editor.moveCursor(QTextCursor.End)

        # Keep the newly focused cell visible in the notebook.
        self.scroll.ensureWidgetVisible(next_cell)


    def run_all(self):
        cells = list(self.cells)

        for cell in cells:
            if cell.editor.toPlainText().strip():
                self.run_cell(cell, advance=False)

        self.statusBar().showMessage("Notebook executed")

    def _rebuild_graph(self):
        # GraphWidget uses the current variable/bar API.
        self._refresh_graph_variables()

    def new_notebook(self):
        self.core = AbacusCore()
        self.core.get_vars().update(self.global_variables)
        self.core.get_funcs().update(self.global_functions)

        for cell in self.cells:
            cell.setParent(None)
            cell.deleteLater()
        self.cells.clear()
        self.graph_points.clear()
        self.graph.set_variables({})
        self.variables_list.clear()
        self.functions_list.clear()
        self._refresh_graph_variables()
        self.add_cell()
        self._refresh_sidebars()
        self.statusBar().showMessage("New notebook")

    def export_notebook(self):
        path, _ = QFileDialog.getSaveFileName(
            self, "Export Notebook", "abacus_notebook.json",
            "Abacus Notebook (*.json)",
        )
        if not path:
            return

        data = {
            "format": "abacus-notebook",
            "version": 1,
            "variables": self.core.get_vars(),
            "functions": self.core.get_funcs(),
            "cells": [c.to_dict() for c in self.cells],
            "graph": [
                {"step": step, "result": value}
                for step, value in self.graph_points
            ],
        }

        try:
            with open(path, "w", encoding="utf-8") as file:
                json.dump(data, file, ensure_ascii=False, indent=2)
        except OSError as exc:
            self.statusBar().showMessage(f"Export error: {exc}")
            return

        self.statusBar().showMessage(f"Notebook exported: {path}")

    def import_notebook(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "Import Notebook", "", "Abacus Notebook (*.json)",
        )
        if not path:
            return

        try:
            with open(path, "r", encoding="utf-8") as file:
                data = json.load(file)

            if data.get("format") != "abacus-notebook":
                raise ValueError("Not an Abacus Notebook file")
            if data.get("version") != 1:
                raise ValueError("Unsupported notebook version")

            variables = data.get("variables", {})
            functions = data.get("functions", {})
            cells = data.get("cells", [])

            if not isinstance(variables, dict) or not isinstance(functions, dict):
                raise ValueError("Invalid variables/functions")
            if not isinstance(cells, list):
                raise ValueError("Invalid cells")

        except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc:
            self.statusBar().showMessage(f"Import error: {exc}")
            return

        self.core = AbacusCore()
        self.core.get_vars().clear()
        self.core.get_vars().update(variables)
        self.core.get_funcs().clear()
        self.core.get_funcs().update(functions)

        for cell in self.cells:
            cell.setParent(None)
            cell.deleteLater()
        self.cells.clear()

        for item in cells:
            cell = NotebookCell(self, code=str(item.get("code", "")))
            if item.get("error"):
                cell.show_error(item["error"])
            elif item.get("result") is not None:
                cell.show_result(item["result"])
            self.cells.append(cell)
            self.notebook_layout.insertWidget(len(self.cells) - 1, cell)

        if not self.cells:
            self.add_cell()

        self._renumber()
        self._update_completion()
        self._refresh_sidebars()
        self._refresh_graph_variables()
        self._rebuild_graph()
        self.statusBar().showMessage(f"Notebook imported: {path}")


def main():
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    app.setStyleSheet(STYLE)

    icon_path = Path(__file__).resolve().parent / "assets" / "abacus_icon.ico"
    if icon_path.exists():
        app.setWindowIcon(QIcon(str(icon_path)))
    window = AbacusNotebook()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()