#!/usr/bin/env python3
import os
import subprocess
import sys
from PyQt6 import QtWidgets, QtGui, QtCore

PRIMARY_SHARED = "/home/primary_shared"

class FolderTree(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Add to Primary Shared")
        self.resize(700, 500)

        layout = QtWidgets.QVBoxLayout(self)

        # buttons
        btn_layout = QtWidgets.QHBoxLayout()
        layout.addLayout(btn_layout)

        self.select_btn = QtWidgets.QPushButton("Select Folder")
        self.select_btn.clicked.connect(self.choose_folder)
        btn_layout.addWidget(self.select_btn)

        self.add_btn = QtWidgets.QPushButton("Add Folder")
        self.add_btn.clicked.connect(self.add_folder)
        btn_layout.addWidget(self.add_btn)

        self.refresh_btn = QtWidgets.QPushButton("Refresh")
        self.refresh_btn.clicked.connect(self.refresh_tree)
        btn_layout.addWidget(self.refresh_btn)

        # tree view
        self.tree = QtWidgets.QTreeWidget()
        self.tree.setHeaderLabels(["Folder"])
        layout.addWidget(self.tree)

        self.selected_folder = None
        self.refresh_tree()

    def choose_folder(self):
        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "Select Folder", os.path.expanduser("~"))
        if folder:
            self.selected_folder = folder
            self.select_btn.setText(f"Selected: {folder}")

    def add_folder(self):
        if not self.selected_folder:
            return
        try:
            subprocess.run(["/usr/local/bin/add_to_primary_shared", self.selected_folder], check=True)
            self.refresh_tree()
            QtWidgets.QMessageBox.information(self, "Success", f"Added {self.selected_folder}")
        except subprocess.CalledProcessError:
            QtWidgets.QMessageBox.critical(self, "Error", f"Failed to add {self.selected_folder}")

    def refresh_tree(self):
        self.tree.clear()
        home = os.path.expanduser("~")
        for root, dirs, files in os.walk(home):
            for d in dirs:
                full_path = os.path.join(root, d)

                # skip symlinks and mounted folders
                if os.path.islink(full_path) or os.path.ismount(full_path):
                    continue

                rel_path = os.path.relpath(full_path, home)
                shared_path = os.path.join(PRIMARY_SHARED, rel_path)

                item = QtWidgets.QTreeWidgetItem([rel_path])

                # determine color
                if os.path.exists(shared_path):
                    # fully moved
                    item.setBackground(0, QtGui.QBrush(QtGui.QColor("green")))
                else:
                    # check if partially moved (any subfolder exists in shared)
                    partially = False
                    try:
                        for f in os.listdir(full_path):
                            f_path = os.path.join(full_path, f)
                            if os.path.islink(f_path) or os.path.ismount(f_path):
                                continue
                            shared_sub = os.path.join(PRIMARY_SHARED, os.path.relpath(f_path, home))
                            if os.path.exists(shared_sub):
                                partially = True
                                break
                    except OSError:
                        # skip folders that cannot be read
                        partially = False

                    if partially:
                        item.setBackground(0, QtGui.QBrush(QtGui.QColor("yellow")))

                self.tree.addTopLevelItem(item)



def main():
    app = QtWidgets.QApplication(sys.argv)
    gui = FolderTree()
    gui.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()