from PyQt4.QtCore import *
from PyQt4.QtGui import *
from c3 import Visitor, astnodes

class AstModelBuilder:
    def __init__(self):
      self.functionIco = QIcon(QPixmap('icons/functionicon.png').scaled(32, 32))
      self.variableIco = QIcon(QPixmap('icons/variableicon.png').scaled(32, 32))
      self.model = QStandardItemModel()
      self.model.setHorizontalHeaderLabels(['Object', 'Type'])

    def build(self, pkg):
        #self.model.clear()
        c = self.model.rowCount()
        self.model.removeRows(0, c)
        self.curItem = self.model.invisibleRootItem()
        if pkg:
            visitor = Visitor()
            visitor.visit(pkg, self.p1, self.p2)

    def p1(self, node):
      if type(node) is astnodes.Variable:
         i = QStandardItem(self.variableIco, str(node))
      elif type(node) is astnodes.Function:
         i = QStandardItem(self.functionIco, str(node))
      elif type(node) is astnodes.Package:
         i = QStandardItem(str(node))
      else:
         return
      typ = str(node.typ) if hasattr(node, 'typ') else ''
      ti = QStandardItem(str(typ))
      ti.setData(node)
      i.setData(node)
      self.curItem.appendRow([i, ti])
      self.curItem = i

    def p2(self, node):
      if type(node) in [astnodes.Variable, astnodes.Function, astnodes.Package]:
         self.curItem = self.curItem.parent()

# The actual widget:
class AstViewer(QTreeView):
    sigNodeSelected = pyqtSignal(object)
    def __init__(self, parent=None):
      super(AstViewer, self).__init__(parent)
      self.clicked.connect(self.selectHandler)
      self.modelBuilder = AstModelBuilder()
      self.setModel(self.modelBuilder.model)

    def setAst(self, ast):
        """ Create a new model and add all ast elements to it """
        self.modelBuilder.build(ast)
        self.expandAll()

    def selectHandler(self, index):
      if not index.isValid():
         return
      model = self.model()
      item = model.itemFromIndex(index)
      node = item.data()
      self.sigNodeSelected.emit(node)