cowork-local/presentation/graph/graph_renderer.py

"""GraphRenderer — the toolbar, scan/render pipeline, and graph/messages
stack of GraphRAG (R08-T14, extracted from
``ui/structure_graph_view.py::StructureGraphView``, lines 188-286/336-661/
664-702 of the original 1035-line file — everything except the right-side
Q&A panel, which is ``graph_qa_widget.py::GraphQaWidget``).

Talks to the Q&A panel only through signals (:attr:`node_selected`,
:attr:`graph_rendered`) and a small read API (:attr:`graph`,
:meth:`selected_node_data`, :attr:`active_project_id`) — this class has no
idea ``GraphQaWidget`` exists, matching how
``presentation/folder/document_preview_manager.py`` doesn't know about the
AI-edit panel either.
"""
from __future__ import annotations

import math
from pathlib import Path
from typing import List, Optional

from PySide6.QtCore import QPointF, Qt, QTimer, QUrl, Signal
from PySide6.QtGui import QColor
from PySide6.QtWidgets import (
    QComboBox, QFileDialog, QGraphicsScene, QHBoxLayout, QLineEdit,
    QPushButton, QStackedWidget, QTabBar, QVBoxLayout, QWidget,
)

from cowork_local.core.worker import AgentWorker
from cowork_local.i18n import on_language_changed, tr
from cowork_local.presentation.graph.graph_messages_view import GraphMessagesView
from cowork_local.presentation.graph.graph_scene_items import _Bridge, _Edge, _GraphView, _Node
from cowork_local.presentation.shared import HAS_WEB_ENGINE
from cowork_local.state import AppContext
from cowork_local.theme import current_palette
from cowork_local.ui.icons import icon


class GraphRenderer(QWidget):
    status_message = Signal(str)
    node_selected = Signal(object)   # a node's .data, whenever the scene selection changes
    graph_rendered = Signal()        # a scan just finished rendering (fresh OR re-fit)
    raw_json_ready = Signal(str)     # pre-formatted HTML for a clicked Messages entry
    project_changed = Signal()       # a DIFFERENT project was selected (or cleared)

    def __init__(self, ctx: AppContext):
        super().__init__()
        self.ctx = ctx
        self._worker: Optional[AgentWorker] = None
        self._node_items: List[_Node] = []
        self._edge_items: List[_Edge] = []
        self._centroid = QPointF(0, 0)
        self._graph = None
        self._needs_scan = False
        self._scan_seq = 0   # only the latest scan's result is rendered (no stale overwrite)
        self._active_project_id = ""  # "" = free path; set = scan locked to that project's sandbox

        self._rescan_timer = QTimer(self)
        self._rescan_timer.setSingleShot(True)
        self._rescan_timer.setInterval(1500)
        self._rescan_timer.timeout.connect(self._scan)

        root = QVBoxLayout(self)
        root.setContentsMargins(0, 0, 0, 0)

        bar = QHBoxLayout()
        self.path_edit = QLineEdit(str(ctx.config.cowork_output_dir()))
        self.path_edit.setPlaceholderText(tr("structure.path_placeholder"))
        self._pick_btn = QPushButton()
        self._pick_btn.setIcon(icon("folder"))
        self._pick_btn.setObjectName("primary")
        self._pick_btn.clicked.connect(self._pick)
        self.project_combo = QComboBox()
        self.project_combo.currentIndexChanged.connect(self._on_project_changed)
        self._scan_btn = QPushButton()
        self._scan_btn.setIcon(icon("search"))
        self._scan_btn.setObjectName("primary")
        self._scan_btn.clicked.connect(self._scan)
        self._export_btn = QPushButton()
        self._export_btn.setIcon(icon("upload"))
        self._export_btn.setObjectName("primary")
        self._export_btn.clicked.connect(self._export)
        bar.addWidget(self.path_edit, 1)
        bar.addWidget(self._pick_btn)
        bar.addWidget(self.project_combo)
        bar.addWidget(self._scan_btn)
        bar.addWidget(self._export_btn)
        root.addLayout(bar)
        self._refresh_project_combo()

        # Đồ thị | Tin nhắn as a real pair of tabs.
        self.view_tabs = QTabBar()
        self.view_tabs.setObjectName("viewTabs")
        self.view_tabs.setDrawBase(False)
        self.view_tabs.setExpanding(False)
        self.view_tabs.addTab(icon("graph"), "")
        self.view_tabs.addTab(icon("message"), "")
        self.view_tabs.currentChanged.connect(self._on_view_tab)
        tab_row = QHBoxLayout()
        tab_row.setContentsMargins(0, 0, 0, 0)
        tab_row.addWidget(self.view_tabs)
        tab_row.addStretch(1)
        root.addLayout(tab_row)

        self.scene = QGraphicsScene()
        self.scene.setBackgroundBrush(QColor(current_palette().bg))
        self.scene.selectionChanged.connect(self._on_selection)
        self.view = _GraphView(self.scene)

        self._stack = QStackedWidget()
        self._stack.addWidget(self.view)
        self.web = None
        self._bridge = None
        self._channel = None
        root.addWidget(self._stack, 1)
        # "Messages" view: all conversation messages grouped BY DAY, shown as
        # JSON — a separate concern composed in (see graph_messages_view.py).
        self._messages = GraphMessagesView(self)

        on_language_changed(self._retranslate)

    def _retranslate(self) -> None:
        self.path_edit.setPlaceholderText(tr("structure.path_placeholder"))
        self._pick_btn.setText(tr("structure.browse"))
        self._scan_btn.setText(tr("structure.scan"))
        self._export_btn.setText(tr("structure.export_png"))
        self.view_tabs.setTabText(0, tr("structure.graph_btn"))
        self.view_tabs.setTabText(1, tr("structure.msgs_btn"))
        self.view_tabs.setTabToolTip(1, tr("structure.msgs_tooltip"))
        self.project_combo.setToolTip(tr("structure.project_tooltip"))
        self._refresh_project_combo()

    # ---- public read API for GraphQaWidget ----------------------------------- #
    @property
    def graph(self):
        return self._graph

    @property
    def active_project_id(self) -> str:
        return self._active_project_id

    def selected_node_data(self) -> list:
        return [item.data for item in self.scene.selectedItems() if isinstance(item, _Node)]

    # ---- project sandbox lock ------------------------------------------------- #
    def _refresh_project_combo(self) -> None:
        from cowork_local.core.projects import list_projects

        keep = self._active_project_id
        self.project_combo.blockSignals(True)
        self.project_combo.clear()
        self.project_combo.addItem(tr("structure.project_none"), "")
        row_to_select = 0
        for i, p in enumerate(list_projects(), start=1):
            self.project_combo.addItem(p.name, p.project_id)
            if p.project_id == keep:
                row_to_select = i
        self.project_combo.setCurrentIndex(row_to_select)
        self.project_combo.blockSignals(False)

    def set_project(self, project_id: str) -> None:
        pid = project_id or ""
        self._refresh_project_combo()
        target = self.project_combo.findData(pid)
        if target < 0:
            target = 0
        if self.project_combo.currentIndex() == target:
            self._on_project_changed(target)
        else:
            self.project_combo.setCurrentIndex(target)

    def _on_project_changed(self, _idx: int) -> None:
        from cowork_local.core.projects import load_project

        pid = self.project_combo.currentData() or ""
        project_changed = pid != self._active_project_id
        self._active_project_id = pid
        locked = bool(pid)
        self.path_edit.setReadOnly(locked)
        self._pick_btn.setEnabled(not locked)
        if locked:
            project = load_project(pid)
            if project is not None:
                self.path_edit.setText(str(project.workspace_dir()))
        if project_changed:
            self.project_changed.emit()   # GraphQaWidget drops its temp extraction cache
            # Mark it and scan on the next visit rather than now — see
            # auto_scan_and_fit()'s docstring for why.
            self._needs_scan = True

    # ---- helpers ---------------------------------------------------------------- #
    def _pick(self) -> None:
        chosen = QFileDialog.getExistingDirectory(self, tr("structure.pick_folder_title"), self.path_edit.text())
        if chosen:
            self.path_edit.setText(chosen)

    def schedule_rescan(self, path: str = "") -> None:
        if self._graph is None:
            self._needs_scan = True
            return
        self._rescan_timer.start()

    # ---- Messages (by day, as JSON) — see graph_messages_view.py --------------- #
    def _on_view_tab(self, index: int) -> None:
        self._messages.on_view_tab(index)

    def _toggle_messages(self) -> None:
        """Kept for callers that still ask for a flip (e.g. keyboard paths)."""
        self._messages.toggle()

    # ---- prewarm / scan lifecycle -------------------------------------------------- #
    def prewarm(self) -> None:
        """Pay for the graph view before it is clicked on, not during."""
        if not HAS_WEB_ENGINE or self.web is not None:
            return
        self._ensure_web()
        if self._graph is None and self.path_edit.text().strip():
            self._needs_scan = False
            self._scan()

    def _ensure_web(self) -> None:
        if self.web is not None or not HAS_WEB_ENGINE:
            return
        from PySide6.QtWebChannel import QWebChannel
        from PySide6.QtWebEngineWidgets import QWebEngineView

        self.web = QWebEngineView()
        self.web.setHtml(
            f"<body style='margin:0;background:{current_palette().bg}'></body>")
        self._bridge = _Bridge()
        self._channel = QWebChannel()
        self._channel.registerObject("py", self._bridge)
        self.web.page().setWebChannel(self._channel)
        self._stack.addWidget(self.web)
        self._stack.setCurrentWidget(self.web)
        if self._graph is not None:
            self._render_d3()

    def auto_scan_and_fit(self) -> None:
        self._ensure_web()
        if not self.path_edit.text().strip():
            return
        if self._worker is not None and self._worker.isRunning():
            self._fit()
            self.graph_rendered.emit()
            return
        if self._graph is not None and not self._needs_scan:
            self._fit()
            self.graph_rendered.emit()
            return
        self._needs_scan = False
        self._scan()

    # ---- scan --------------------------------------------------------------------- #
    def _scan(self) -> None:
        path = self.path_edit.text().strip() or str(Path.cwd())
        mode = "files"
        use_cmem = bool(self.ctx.config.codebase_memory.get("enabled"))
        cmem_bin = self.ctx.config.codebase_memory.get("binary_path", "")
        st = self.ctx.config.structure
        max_nodes = int(st.get("max_nodes", 500) or 0)
        max_edges = int(st.get("max_edges", 500) or 0)
        self._scan_seq += 1
        seq = self._scan_seq
        self.status_message.emit(tr("structure.scanning"))

        def job(worker: AgentWorker):
            from cowork_local.core.structure_graph import (
                build_from_codebase_memory, build_from_directory, force_layout,
            )
            if use_cmem:
                from cowork_local.core.codebase_memory import CodebaseMemory
                mem = CodebaseMemory(cmem_bin)
                graph = (build_from_codebase_memory(mem, path, mode, max_nodes, max_edges)
                         if mem.available else build_from_directory(path, mode, max_nodes, max_edges))
            else:
                graph = build_from_directory(path, mode, max_nodes, max_edges)
            pos = force_layout(graph)
            return {"graph": graph, "pos": pos, "seq": seq}

        w = AgentWorker(job)
        w.finished_ok.connect(self._render)
        w.failed.connect(lambda e: self.status_message.emit(tr("structure.scan_error", err=e)))
        self._worker = w
        w.start()

    def _render(self, result: dict) -> None:
        if result.get("seq") is not None and result["seq"] != self._scan_seq:
            return
        graph = result.get("graph")
        pos = result.get("pos", {})
        if graph is None:
            return
        self._graph = graph

        self.scene.clear()
        self.scene.setBackgroundBrush(QColor(current_palette().bg))
        self._node_items = []
        self._edge_items = []
        degree = {n.id: 0 for n in graph.nodes}
        for e in graph.edges:
            if e.source in degree:
                degree[e.source] += 1
            if e.target in degree:
                degree[e.target] += 1
        items = {}
        sx = sy = 0.0
        for node in graph.nodes:
            radius = int(8 + min(20, 2.2 * math.sqrt(degree.get(node.id, 0))))
            item = _Node(node, radius)
            x, y = pos.get(node.id, (0, 0))
            item.setPos(x, y)
            self.scene.addItem(item)
            items[node.id] = item
            self._node_items.append(item)
            sx += x
            sy += y
        for edge in graph.edges:
            a, b = items.get(edge.source), items.get(edge.target)
            if a and b:
                e = _Edge(a, b, getattr(edge, "type", ""))
                self.scene.addItem(e)
                self._edge_items.append(e)
        n = max(1, len(self._node_items))
        self._centroid = QPointF(sx / n, sy / n)
        self._fit()

        if self.web is not None:
            self._render_d3()

        note = tr("structure.truncated_note") if getattr(graph, "truncated", False) else ""
        self.status_message.emit(tr(
            "structure.graph_summary", nodes=len(graph.nodes), edges=len(graph.edges), note=note))
        self.graph_rendered.emit()

    def _render_d3(self) -> None:
        if self.web is None or self._graph is None:
            return
        from cowork_local.core.d3_graph import build_html
        try:
            self.web.setHtml(build_html(self._graph), QUrl("https://cowork.local/"))
        except Exception as exc:
            self.status_message.emit(f"D3 view error: {exc}")

    # ---- native interactions ------------------------------------------------------- #
    def _on_selection(self) -> None:
        for item in self.scene.selectedItems():
            if isinstance(item, _Node):
                self.node_selected.emit(item.data)
                return

    def _fit(self) -> None:
        if self.web is not None and self._stack.currentWidget() is self.web:
            self.web.page().runJavaScript("window.fitGraph && window.fitGraph();")
            return
        rect = self.scene.itemsBoundingRect()
        if not rect.isNull():
            self.view.fitInView(rect.adjusted(-40, -40, 40, 40), Qt.KeepAspectRatio)

    def _export(self) -> None:
        path, _ = QFileDialog.getSaveFileName(
            self, tr("structure.export_title"), "structure-graph.png", "PNG (*.png)")
        if not path:
            return
        showing_d3 = (self.web is not None and self._stack.currentWidget() is self.web)
        if showing_d3:
            self._export_d3_png(path)
        else:
            self._export_widget_grab(path)

    def _export_d3_png(self, path: str) -> None:
        def on_result(data_url) -> None:
            if not isinstance(data_url, str) or "," not in data_url:
                self._export_widget_grab(path)
                return
            import base64
            try:
                with open(path, "wb") as f:
                    f.write(base64.b64decode(data_url.split(",", 1)[1]))
                self.status_message.emit(tr("structure.export_done", path=path))
            except (OSError, ValueError) as exc:
                self.status_message.emit(tr("structure.export_failed", err=str(exc)))
        self.web.page().runJavaScript("window.exportPng ? window.exportPng() : ''", on_result)

    def _export_widget_grab(self, path: str) -> None:
        ok = self._stack.currentWidget().grab().save(path, "PNG")
        if ok:
            self.status_message.emit(tr("structure.export_done", path=path))
        else:
            self.status_message.emit(tr("structure.export_failed", err="grab() returned no image"))


__all__ = ["GraphRenderer"]