2 files changed, 31 insertions, 0 deletions

diff --git a/src/Sequence/Contact/Types.hs b/src/Sequence/Contact/Types.hs
index dff886d..47687b7 100644
--- a/src/Sequence/Contact/Types.hs
+++ b/src/Sequence/Contact/Types.hs
@@ -36,6 +36,8 @@ hitzone = iso (Hitzone . CI.mk) (CI.original . _hitzone)
 instance IsString Hitzone where
   fromString = view hitzone
 
+makePrisms ''Hitzone
+
 makeLenses ''Armor
 
 instance Default Armor where
diff --git a/src/Sequence/Utils.hs b/src/Sequence/Utils.hs
index 517c3c2..bbd1477 100644
--- a/src/Sequence/Utils.hs
+++ b/src/Sequence/Utils.hs
@@ -11,6 +11,7 @@ module Sequence.Utils
 import Sequence.Types
 
 import Control.Monad.State.Strict
+import Control.Monad.Trans.Maybe
 
 import Control.Applicative
 import Control.Monad
@@ -25,6 +26,9 @@ import qualified Data.CaseInsensitive as CI
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 
+import Data.Set (Set)
+import qualified Data.Set as Set
+
 import Data.Function
 import Data.Default
 import Data.Maybe
@@ -32,12 +36,14 @@ import Text.Read (readMaybe)
 
 import Data.List
 import Data.Bool
+import Data.Char
 
 import System.Console.Shell
 import System.Console.Shell.ShellMonad
 import System.Console.Shell.Backend.Haskeline
 
 import Sequence.Utils.Ask
+import Sequence.Contact.Types
 
 class Argument a st | a -> st where
   arg :: String -> Sh st (Maybe a)
@@ -85,3 +91,26 @@ instance Completion Faction GameState where
 
 instance Argument Faction GameState where
   arg = return . Just . flip (set faction') def
+
+instance Completion (Set Hitzone) GameState where
+  completableLabel _ = "<hitzones>"
+  complete _ st (over each reverse . span (/= ',') . reverse -> (wPrefix, lPrefix))
+    | not $ null wPrefix
+    , all (== '*') wPrefix = return . pure . join $ hitzones
+    | otherwise = return . map (lPrefix ++) . filter ((isPrefixOf `on` CI.foldCase) wPrefix) $ hitzones
+    where
+      hitzones = sort . map (review hitzone) $ fromMaybe [] (Map.keys <$> preview (gFocus' . eStats . sHitzones) st)
+      join [] = ""
+      join [x] = x
+      join (x:xs) = x ++ "," ++ join xs
+
+instance Argument (Set Hitzone) GameState where
+  arg protoWs = runMaybeT $ do
+    let
+      trim = reverse . dropWhile isSpace . reverse . dropWhile isSpace
+      split = foldr (\c l@(w:ws) -> if c == ',' then "" : l else (c : w) : ws) [""]
+      ws = Set.fromList . map CI.mk . filter (not . null) . map trim . split $ protoWs
+      hasGlob = Set.member "*" ws
+    hitzones <- Set.map (view _Hitzone) . Map.keysSet <$> MaybeT (preuse $ gFocus' . eStats . sHitzones)
+    guard (hasGlob || hitzones `Set.isSubsetOf` ws)
+    return . Set.map (review _Hitzone) $ if hasGlob then hitzones else ws