From d4632b90598d709bbf3e39e38f3e8d3f510cebb1 Mon Sep 17 00:00:00 2001
From: blitzmann <ryan.xgamer99@gmail.com>
Date: Sun, 20 Mar 2016 22:28:18 -0400
Subject: [PATCH] Add context menu to apply implant set to fit

---
 gui/builtinContextMenus/__init__.py    |  1 +
 gui/builtinContextMenus/implantSets.py | 53 ++++++++++++++++++++++++++
 gui/implantView.py                     | 17 ++++++++-
 3 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 gui/builtinContextMenus/implantSets.py

diff --git a/gui/builtinContextMenus/__init__.py b/gui/builtinContextMenus/__init__.py
index 785e3eebf..c2a39b179 100644
--- a/gui/builtinContextMenus/__init__.py
+++ b/gui/builtinContextMenus/__init__.py
@@ -20,4 +20,5 @@ __all__ = [
     "priceClear",
     "amount",
     "metaSwap",
+    "implantSets",
 ]
diff --git a/gui/builtinContextMenus/implantSets.py b/gui/builtinContextMenus/implantSets.py
new file mode 100644
index 000000000..c96080983
--- /dev/null
+++ b/gui/builtinContextMenus/implantSets.py
@@ -0,0 +1,53 @@
+from gui.contextMenu import ContextMenu
+from gui.itemStats import ItemStatsDialog
+import eos.types
+import gui.mainFrame
+import service
+import gui.globalEvents as GE
+import wx
+
+class ImplantSets(ContextMenu):
+    def __init__(self):
+        self.mainFrame = gui.mainFrame.MainFrame.getInstance()
+
+    def display(self, srcContext, selection):
+        return srcContext == "implantView"
+
+    def getText(self, itmContext, selection):
+        return "Add Implant Set"
+
+    def getSubMenu(self, context, selection, rootMenu, i, pitem):
+        m = wx.Menu()
+        bindmenu = rootMenu if "wxMSW" in wx.PlatformInfo else m
+
+        sIS = service.ImplantSets.getInstance()
+        implantSets = sIS.getImplantSetList()
+
+        self.idmap = {}
+
+        for set in implantSets:
+            id = wx.NewId()
+            mitem = wx.MenuItem(rootMenu, id, set.name)
+            bindmenu.Bind(wx.EVT_MENU, self.handleSelection, mitem)
+            self.idmap[id] = set
+            m.AppendItem(mitem)
+
+        return m
+
+    def handleSelection(self, event):
+        set = self.idmap.get(event.Id, None)
+
+        if set is None:
+            event.Skip()
+            return
+
+        sFit = service.Fit.getInstance()
+        fitID = self.mainFrame.getActiveFit()
+        for implant in set.implants:
+            print implant.item.ID, implant.item.name
+            sFit.addImplant(fitID, implant.item.ID)
+
+        wx.PostEvent(self.mainFrame, GE.FitChanged(fitID=fitID))
+
+
+ImplantSets.register()
diff --git a/gui/implantView.py b/gui/implantView.py
index 5c065094a..4b4755252 100644
--- a/gui/implantView.py
+++ b/gui/implantView.py
@@ -172,9 +172,15 @@ class ImplantDisplay(d.Display):
 
     def spawnMenu(self):
         sel = self.GetFirstSelected()
+        menu = None
+
+        sFit = service.Fit.getInstance()
+        fit = sFit.getFit(self.mainFrame.getActiveFit())
+
+        if not fit:
+            return
+
         if sel != -1:
-            sFit = service.Fit.getInstance()
-            fit = sFit.getFit(self.mainFrame.getActiveFit())
             implant = fit.appliedImplants[sel]
 
             sMkt = service.Market.getInstance()
@@ -182,4 +188,11 @@ class ImplantDisplay(d.Display):
             itemContext = sMkt.getCategoryByItem(implant.item).name
 
             menu = ContextMenu.getMenu((implant,), (sourceContext, itemContext))
+        elif sel == -1 and fit.implantSource == ImplantLocation.FIT:
+            fitID = self.mainFrame.getActiveFit()
+            if fitID is None:
+                return
+            context = (("implantView",),)
+            menu = ContextMenu.getMenu([], *context)
+        if menu is not None:
             self.PopupMenu(menu)