#####################################
# Copyright (c) openBrain 2020
# Licensed under LGPL v2
#
# This FreeCAD macro prints the Tree of selected object(s) as ASCII art.
# Several styles are available. Tree branch pattern is customizable.
# It can export to clipboard, file or embedded text document.
# This mainly aims at documenting the model.
#
# Version history:
# *0.7: add option to exclude origin objects
# *0.6: introduce last fork string definition, add new style
# *0.5: beta release
#
#####################################

from __future__ import annotations


__Name__ = 'TreeToAscii'
__Comment__ = 'Prints the Tree of selected object(s) as ASCII art'
__Author__ = 'openBrain'
__Version__ = '0.7.1'
__Date__ = '2020-03-23'
__License__ = 'LGPL v2'
__Web__ = 'https://www.freecadweb.org/wiki/Macro_TreeToAscii'
__Wiki__ = 'https://www.freecadweb.org/wiki/Macro_TreeToAscii'
__Icon__ = ''
__Help__ = 'Select object(s) you want to print the tree then run'
__Status__ = 'Beta'
__Requires__ = 'FreeCAD >= 0.17'
__Communication__ = 'https://forum.freecadweb.org/viewtopic.php?f=22&t=43504'

import FreeCAD as app
import FreeCADGui as gui
from PySide import QtGui, QtCore  # Special FreeCAD's PySide, which may be PySide2!

PREF_PATH = 'User parameter:BaseApp/Preferences/Mod/TreeToAscii'
PREF_PATTERN = 'TreePattern'
PREF_STYLE = 'TreeStyle'
# Define styles as NAME: [branch str, fork str, opt:last fork str].
STYLES = {
    'Modern': [' │ ', ' ├ ', ' └ '],  # U+2502, U+251C, U+2502
    'Baroque': [' │  ', ' ├─ ',  ' └─ '],  # U+2502, U+251C + U+2500, U+2514 + U+2500
    'ASCII': [' | ', ' \\_']
    }
DEF_PATTERN = '{LBL} ({SNAM})'
DEF_STYLE = list(STYLES.keys())[0]

g_ascii_tree = ''


def print_children(objs: list, top_level=True, baseline='', excluded=None):
    """Build ASCII tree by recursive call.

    Parameters
    ==========
    objs: List objects to print and start the recursion from.
    baseline: String to print before the object name.
    excluded: List of objects to exclude from the tree.
              As of v0.7.0, optionaly exclude part origins and origin's children.

    """
    global g_ascii_tree
    if excluded is None:
        excluded = []
    for i, obj in enumerate(objs, start=1):
        data = {
            'LBL': obj.Label,
            'NAM': obj.Name,
            'SNAM': ''.join(filter(str.isalpha, obj.Name)),
            'TYP': obj.TypeId,
            'BTYP': obj.TypeId.partition('::')[0],
            'STYP': obj.TypeId.rpartition('::')[-1]
        }
        if (
                getattr(obj, 'TypeId', '') == 'App::Part'
                and not gen_include_origins.isChecked()
        ):
            excluded.append(obj.Origin)
            excluded += obj.Origin.OriginFeatures

        if obj in excluded:
            continue
        branch_str = STYLES[gen_style.currentText()][0]
        fork_str = STYLES[gen_style.currentText()][1]
        last_fork_str = STYLES[gen_style.currentText()][-1]
        if i == len(objs):
            # The last object.
            if not top_level:
                g_ascii_tree += baseline + last_fork_str
        else:
            if not top_level:
                g_ascii_tree += baseline + fork_str
        try:
            g_ascii_tree += gen_pattern.text().format(**data)
        except (ValueError, IndexError, KeyError):
            g_ascii_tree = 'Error in generating pattern'
            return
        g_ascii_tree += '\n'
        if i == len(objs):
            # The last object.
            if top_level:
                baseline_children = ''
            else:
                baseline_children = baseline + ' ' * len(branch_str)
        else:
            if top_level:
                baseline_children = ''
            else:
                baseline_children = baseline + branch_str
        print_children(
                obj.ViewObject.claimChildren(),
                False,
                baseline_children,
                excluded,
        )


def copy_to_clipboard():
    """Copy the global ASCII tree to clipboard."""
    tree_view.selectAll()
    tree_view.copy()
    tree_view.setPlainText('Copied to clipboard')
    QtCore.QTimer.singleShot(3000, lambda:tree_view.setPlainText(g_ascii_tree))


def save_to_file():
    """Save the global ASCII tree to file system."""
    path = QtGui.QFileDialog.getSaveFileName()
    if path[0]:
        filename = path[0]
        try:
            with open(filename, 'w') as f:
                f.write(g_ascii_tree)
        except OSError:
            tree_view.setPlainText(f'ERROR: Can not open file: "{filename}"')
        else:
            tree_view.setPlainText(f'Saved to file: "{filename}"')
        QtCore.QTimer.singleShot(3000, lambda:tree_view.setPlainText(g_ascii_tree))


def embed_in_doc():
    """Store the global ASCII tree as `App::TextDocument` object."""
    try:
        text_obj = app.ActiveDocument.addObject('App::TextDocument', 'ModelTree')
    except TypeError:
        tree_view.setPlainText('ERROR: can not create a Text object. Maybe your FreeCAD version is too old')
        return
    text_obj.Text = g_ascii_tree
    tree_view.setPlainText(f'Embedded as text object "{text_obj.Label}"')
    QtCore.QTimer.singleShot(3000, lambda: tree_view.setPlainText(g_ascii_tree))


def update_text():
    """Update the global ASCII tree."""
    global g_ascii_tree
    g_ascii_tree = ''
    print_children(gui.Selection.getSelection())
    tree_view.setPlainText(g_ascii_tree)


def on_generate():
    """Called when the 'Generate' button is clicked."""
    p.SetString(PREF_PATTERN, gen_pattern.text())
    p.SetString(PREF_STYLE, gen_style.currentText())
    update_text()


# Get saved parameters or use defaults.
p = app.ParamGet(PREF_PATH)
g_name_pattern = p.GetString(PREF_PATTERN, DEF_PATTERN)
g_tree_style = p.GetString(PREF_STYLE, DEF_STYLE)
if g_tree_style not in STYLES:
    g_tree_style = DEF_STYLE


# Build UI and show
dlg = QtGui.QDialog(gui.getMainWindow(), QtCore.Qt.Tool)
dlg_layout = QtGui.QVBoxLayout(dlg)
gen_layout = QtGui.QHBoxLayout()
gen_pattern = QtGui.QLineEdit(g_name_pattern, dlg)
gen_pattern.setToolTip('''Defined variables are:
    {LBL}: object Label
    {NAM}: object Name
    {SNAM}: object ShortName (without digits)
    {TYP}: object Type
    {BTYP}: object BaseType
    {STYP}: object ShortType (without base)
If you need to print curly braces, they have to be doubled''')
gen_run = QtGui.QPushButton('Generate', dlg)
gen_run.clicked.connect(on_generate)
gen_style = QtGui.QComboBox(dlg)
gen_style.addItems(list(STYLES.keys()))
gen_style.setCurrentIndex(list(STYLES.keys()).index(g_tree_style))
gen_include_origins = QtGui.QCheckBox('Include origins', dlg)
gen_include_origins.setChecked(True)
gen_include_origins.stateChanged.connect(on_generate)
gen_layout.addWidget(QtGui.QLabel('Generating pattern: '))
gen_layout.addWidget(gen_pattern)
gen_layout.addWidget(QtGui.QLabel('Style: '))
gen_layout.addWidget(gen_style)
gen_layout.addWidget(gen_include_origins)
gen_layout.addWidget(gen_run)
tree_view = QtGui.QPlainTextEdit(g_ascii_tree)
tree_view.setReadOnly(True)
tree_view.setMinimumWidth(gui.getMainWindow().width() // 2)
tree_view.setMinimumHeight(gui.getMainWindow().height() // 2)
tree_view.setLineWrapMode(QtGui.QPlainTextEdit.NoWrap)
button_box = QtGui.QDialogButtonBox()
copy_clip_but = QtGui.QPushButton('Copy to clipboard', button_box)
button_box.addButton(copy_clip_but, QtGui.QDialogButtonBox.ActionRole)
copy_clip_but.clicked.connect(copy_to_clipboard)
_but = QtGui.QPushButton('Save to file', button_box)
button_box.addButton(_but, QtGui.QDialogButtonBox.ActionRole)
_but.clicked.connect(save_to_file)
store_embed_but = QtGui.QPushButton('Embed in document', button_box)
button_box.addButton(store_embed_but, QtGui.QDialogButtonBox.ActionRole)
store_embed_but.clicked.connect(embed_in_doc)
close_dlg_but = QtGui.QPushButton('Close', button_box)
button_box.addButton(close_dlg_but, QtGui.QDialogButtonBox.RejectRole)
button_box.rejected.connect(dlg.reject)
dlg_layout.insertLayout(0, gen_layout)
dlg_layout.addWidget(tree_view)
dlg_layout.addWidget(button_box)
update_text()
dlg.exec_()