content: {-# LANGUAGE UnicodeSyntax, NoImplicitPrelude #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE PatternGuards #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module Tmk where
  
  import BasePrelude hiding (toList)
  import Prelude.Unicode hiding ((∈))
  import Data.Foldable.Unicode ((∈))
  import Data.List.Unicode ((∖))
  import Data.Monoid.Unicode ((⊕))
  import qualified Data.Set.Unicode as S
  import Util (show', toString, (>$>), filterOnFst, groupWith', versionStr)
  import WithBar (WithBar(..))
  import WithPlus (WithPlus(..))
  
  import Control.Monad.State (State, runState, get, gets, modify)
  import Control.Monad.Writer (tell)
  import Data.Foldable (toList)
  import qualified Data.List.NonEmpty as NE
  import Data.Map (Map)
  import qualified Data.Map as M
  import Data.Set (Set)
  import qualified Data.Set as S
  import Lens.Micro.Platform (view, over, makeLenses, _1, _2)
  
  import Layout.Key (Key(..), getLevel, filterKeyOnShiftstatesM)
  import Layout.Layout
  import Layout.Modifier (toBaseModifier, getEqualModifiers)
  import Layout.ModifierEffect (defaultModifierEffect)
  import Layout.Types
  import Lookup.Linux (actionAndRedirect, charAndString)
  import Lookup.Tmk
  import PresetLayout (defaultKeys, defaultFullLayout)
  
  prepareLayout ∷ Logger m ⇒ Layout → m Layout
  prepareLayout =
      addSingletonKeysAsKeys >>>
      addDefaultKeys defaultKeys >>>
      _keys
          ( filterM (supportedPos ∘ view _pos) >=>
            traverse (filterKeyOnShiftstatesM supportedShiftstate)
          )
  
  supportedPos ∷ Logger m ⇒ Pos → m Bool
  supportedPos pos
    | pos S.∈ __usedPosses unimap = pure True
    | otherwise = False <$ tell [show' pos ⊕ " is not supported in TMK"]
  
  supportedShiftstate ∷ Logger m ⇒ Shiftstate → m Bool
  supportedShiftstate = fmap and ∘ traverse supportedModifier ∘ toList
  
  supportedModifier ∷ Logger m ⇒ Modifier → m Bool
  supportedModifier modifier
      | defaultModifierEffect modifier ≡ Shift = pure True
      | otherwise = False <$ tell [show' modifier ⊕ " is not supported for layer selection in TMK"]
  
  data TmkLetter
      = TmkAction String
      | TmkFn String
      | TmkModifier Modifier
      | TmkMacro String [String] [String]
      deriving (Eq, Show, Read)
  
  isTmkModifier ∷ TmkLetter → Bool
  isTmkModifier (TmkModifier _) = True
  isTmkModifier _ = False
  
  toMacro ∷ TmkLetter → Maybe (String, [String], [String])
  toMacro (TmkMacro mId press release) = Just (mId, press, release)
  toMacro _ = Nothing
  
  type Tmk = State [String]
  
  printTmkLetter ∷ Set Modifier → TmkLetter → Tmk String
  printTmkLetter _ (TmkAction s) = pure s
  printTmkLetter _ (TmkFn s) = writeAction s
  printTmkLetter modifiersInfluenced (TmkModifier modifier)
    | modifier ∈ modifiersInfluenced
    = writeAction ("ACTION_FUNCTION_OPT(F_MODIFIER, MOD_" ⊕ map toUpper (toString modifier) ⊕ ")")
    | otherwise = pure $ fromMaybe "NO" (lookup modifier modifierAndKeycode)
  printTmkLetter _ (TmkMacro macroId _ _) = writeAction ("ACTION_MACRO(" ⊕ macroId ⊕ ")")
  
  writeAction ∷ String → Tmk String
  writeAction s = findIndex (≡ s) <$> get >>= maybe
      (gets (("FN" ⊕) ∘ show ∘ length) <* modify (⧺ [s]))
      (pure ∘ ("FN" ⊕) ∘ show)
  
  data TmkLayer = TmkLayer
      { __tmkShiftstate ∷ Shiftstate
      , __tmkLetters ∷ Map Pos TmkLetter
      } deriving (Show, Read)
  makeLenses ''TmkLayer
  
  printTmkLayer ∷ Set Modifier → Int → (Map Pos TmkLetter, Shiftlevel) → Tmk [String]
  printTmkLayer modifiersInfluenced i (tmkLetters, shiftlevel) = do
      actions ← zipWithM (zipWithM printPos) (__sizes unimap) (__posses unimap)
      pure (
          [ "// " ⊕ toString shiftlevel
          , "[" ⊕ show i ⊕ "] = " ⊕ __name unimap ⊕ "("
          ] ⧺ addCommas (map (intercalate ",") actions) ⧺
          [ "),"
          ])
    where
      printPos size =
          printTmkLetter modifiersInfluenced ∘ getTmkLetter >$>
          printf ("%" ⊕ show size ⊕ "s")
      getTmkLetter pos
        | i ≡ 0 = fromMaybe (TmkAction "NO") $ asum
            [ M.lookup pos tmkLetters
            , TmkModifier <$> lookup pos posAndModifier
            , TmkAction <$> lookup pos posAndTmkAction
            ]
        | otherwise = fromMaybe (TmkAction "TRNS") $ M.lookup pos tmkLetters
      addCommas [] = []
      addCommas [x] = [x]
      addCommas (x:xs) = x ⊕ "," : addCommas xs
  
  data TmkKeymap = TmkKeymap
      { __tmkGetMaxIndex ∷ Shiftstate → Int
      , __tmkLayers ∷ [TmkLayer]
      }
  makeLenses ''TmkKeymap
  
  printTmkKeymap ∷ TmkKeymap → String
  printTmkKeymap (TmkKeymap getMaxIndex layers') = unlines $
      [ "// Generated by KLFC " ⊕ versionStr
      , "// https://github.com/39aldo39/klfc"
      , ""
      , "#include \"unimap_trans.h\""
      , "#include \"action_util.h\""
      , "#include \"action_layer.h\""
      , ""
      , "enum function_id {"
      , "    F_MODIFIER,"
      , "};"
      , ""
      ] ⧺ bool (
      [ "enum macro_id {"
      ] ⧺ map (\macro → replicate 4 ' ' ⊕ view _1 macro ⊕ ",") macros ⧺
      [ "};"
      , ""
      ]) [] (null macros) ⧺ bool (
      [ "enum modifier_id {"
      ] ⧺ map (\m → replicate 4 ' ' ⊕ showMod m ⊕ ",") modifiersInfluencedList ⧺
      [ "};"
      , ""
      ]) [] (null modifiersInfluencedList)
      ⧺ zipWith (\i action → "#define AC_FN" ⊕ show i ⊕ " " ⊕ action) [0 ∷ Int ..] functionKeys ⧺
      [ ""
      , "#ifdef KEYMAP_SECTION_ENABLE"
      , "const action_t actionmaps[][UNIMAP_ROWS][UNIMAP_COLS] __attribute__ ((section (\".keymap.keymaps\"))) = {"
      , "#else"
      , "const action_t actionmaps[][UNIMAP_ROWS][UNIMAP_COLS] PROGMEM = {"
      , "#endif"
      ] ⧺ map (replicate 4 ' ' ⊕) printedLayers ⧺
      [ "};"
      , ""
      ] ⧺ bool (
      [ "const macro_t *action_get_macro(keyrecord_t *record, uint8_t id, uint8_t opt) {"
      , "    switch (id) {"
      ] ⧺ concatMap (map (replicate 8 ' ' ⊕) ∘ printMacro) macros ⧺
      [ "    }"
      , "    return MACRO_NONE;"
      , "}"
      , ""
      ]) [] (null macros)
      ⧺ map (\m → "#define " ⊕ showMask m ⊕ " (" ⊕ intercalate "|" (map (\m' → "MOD_BIT(" ⊕ showKeycode m' ⊕ ")") (modifierToRealModifiers m)) ⊕ ")") realModsInShiftstates
      ⧺ bool [""] [] (null realModsInShiftstates)
      ⧺ bool (
      [ "enum vmod_mask {"
      ] ⧺ zipWith (\i m → replicate 4 ' ' ⊕ showMask m ⊕ " = " ⊕ show (2^i ∷ Int) ⊕ ",") [0 ∷ Int ..] virtualModsInShiftstates ⧺
      [ "};"
      , ""
      ]) [] (null virtualModsInShiftstates) ⧺
      [ "uint" ⊕ show indexSize ⊕ "_t vmods = 0;"
      , ""
      , "const uint" ⊕ show layerSize ⊕ "_t layer_states[] = {"
      ] ⧺ map (\(i, state) → replicate 4 ' ' ⊕ "0x" ⊕ showHex (layerStateAfterGrouping i) "," ⊕ " // " ⊕ toString state) (getLayers getMaxIndex layers') ⧺
      [ "};"
      , ""
      , "void action_function(keyrecord_t *record, uint8_t id, uint8_t opt) {"
      , "    uint8_t pressed = record->event.pressed;"
      , "    switch (id) {"
      , "        case F_MODIFIER:"
      , "            // Set the new modifier"
      , "            switch (opt) {"
      ] ⧺
      map (\m → "                case " ⊕ showMod m ⊕ ": pressed ? add_key(" ⊕ showKeycode m ⊕ ") : del_key(" ⊕ showKeycode m ⊕ "); break;") realModsInfluenced ⧺
      map (\m → "                case " ⊕ showMod m ⊕ ": pressed ? (vmods |= " ⊕ showMask m ⊕ ") : (vmods &= ~" ⊕ showMask m ⊕ "); break;") virtualModsInfluenced ⧺
      [ "            }"
      , ""
      , "            // Update the layer"
      , "            uint8_t mods = get_mods();"
      , "            uint" ⊕ show indexSize ⊕ "_t layer_index = 0;"
      ] ⧺ zipWith (\i m → "            layer_index |= mods & " ⊕ showMask m ⊕ " ? " ⊕ show (2^i ∷ Int) ⊕ " : 0;") [0 ∷ Int ..] realModsInShiftstates ⧺
      [ "            layer_index |= vmods << " ⊕ show (length realModsInShiftstates) ⊕ ";"
      , "            layer_clear();"
      , "            layer_or(layer_states[layer_index]);"
      , "            break;"
      , "    }"
      , "}"
      ]
    where
      (printedLayers, functionKeys) = flip runState [] $
          concat <$> zipWithM (printTmkLayer modifiersInfluenced) [0..] layers
      shiftstates = map __tmkShiftstate layers'
      (layerStateAfterGrouping, layers) = groupLayers layers'
      modifiersInShiftstates = S.unions ∘ map getSet $ shiftstates
      modifiersInfluenced = S.fromList ∘ concatMap (concatMap getEqualModifiers ∘ toList) $ shiftstates
      modifiersInfluencedList = toList modifiersInfluenced
      showMod = ("MOD_" ⊕) ∘ map toUpper ∘ toString
      showKeycode modifier =
          let e = error (show' modifier ⊕ " is not a real modifier in TMK")
          in "KC_" ⊕ fromMaybe e (lookup modifier modifierAndKeycode)
      showMask modifier =
          bool "V" "" (isRealModifier modifier) ⊕
          showMod (toBaseModifier modifier) ⊕ "_MASK"
      (realModsInShiftstates, virtualModsInShiftstates) = partition isRealModifier (toList modifiersInShiftstates)
      (realModsInfluenced, virtualModsInfluenced) = partition isRealModifier modifiersInfluencedList
      asPower2 x
        | x < 1     = 8
        | otherwise = max 8 (bit (ceiling (logBase 2 (fromIntegral x ∷ Double)) ∷ Int)) ∷ Int
      indexSize = asPower2 (length modifiersInShiftstates)
      layerSize = asPower2 (length layers')
  
      macros = nub (mapMaybe toMacro (concatMap (M.elems ∘ fst) layers))
      printMacro (mId, press, release) =
          [ "case " ⊕ mId ⊕ ":"
          , "    return record->event.pressed ?"
          , "           " ⊕ printMacroPart press ⊕ " :"
          , "           " ⊕ printMacroPart release ⊕ ";"
          ]
      printMacroPart [] = "MACRO_NONE"
      printMacroPart parts = "MACRO(" ⊕ concatMap (⊕ ", ") parts ⊕ "END)"
  
  getLayers ∷ (Shiftstate → Int) → [TmkLayer] → [(Int, Shiftstate)]
  getLayers getMaxIndex layers = ($ shiftstates) $
      S.toAscList ∘ S.unions ∘ map getSet >>> -- get an ascending list of modifiers
      uncurry (⧺) ∘ partition isRealModifier >>> -- place the real modifiers at the front
      map (WithPlus ∘ S.fromList) ∘ subsequences >>> -- generate the substates with the real modifiers at the front
      map (sum ∘ map bit ∘ getActiveStates getMaxIndex shiftstates &&& id)
    where
      shiftstates = map __tmkShiftstate layers
  
  getActiveStates ∷ (Shiftstate → Int) → [Shiftstate] → Shiftstate → [Int]
  getActiveStates getMaxIndex states state =
      findIndices (`isSubState` state) (take (getMaxIndex state) states)
  
  toTmkKeymap ∷ Logger m ⇒ Layout → m TmkKeymap
  toTmkKeymap =
      prepareLayout >=>
      toTmkKeymap' >$>
      over _tmkLayers addTransLetters >>>
      removeTransLayers
  
  toTmkKeymap' ∷ Logger m ⇒ Layout → m TmkKeymap
  toTmkKeymap' layout =
      TmkKeymap getMaxIndex <$> zipWithM lettersToTmkLayer shiftstates letters
    where
      (keys, shiftstates) = unifyShiftstates (view _keys layout)
      shiftlevels = map (WithBar ∘ (:| [])) shiftstates
      posses = map (view _pos) keys
      letters = map (M.fromList ∘ zip posses) ∘ transpose ∘ map (view _letters) $ keys
      getMaxIndex = maybe (length shiftstates) (succ ∘ fst) ∘ getLevel keyForShiftstates
      keyForShiftstates = Key undefined undefined shiftlevels undefined (Just False)
  
  removeTransLayers ∷ TmkKeymap → TmkKeymap
  removeTransLayers (TmkKeymap getMaxIndex layers) = TmkKeymap getMaxIndex' layers'
    where
      (layers', filteredIndices) = ($ layers) $
          map (mfilter (not ∘ allTrans) ∘ Just) >>>
          catMaybes &&& findIndices isJust
      getMaxIndex' = (\i → length (takeWhile (< i) filteredIndices)) ∘ getMaxIndex
  
  lettersToTmkLayer ∷ Logger m ⇒ Shiftstate → Map Pos Letter → m TmkLayer
  lettersToTmkLayer state = traverse (letterToTmkLetter state) >$> TmkLayer state
  
  letterToTmkLetter ∷ Logger m ⇒ Shiftstate → Letter → m TmkLetter
  letterToTmkLetter state letter = letterToTmkLetter' letter state letter
  
  letterToTmkLetter' ∷ Logger m ⇒ Letter → Shiftstate → Letter → m TmkLetter
  letterToTmkLetter' _ _ LNothing = pure (TmkAction "NO")
  letterToTmkLetter' _ _ (Ligature _ s)
    | Just tmkLetter ← tmkLetterByMacro (map Char s) = pure tmkLetter
  letterToTmkLetter' errorL state (Action a)
    | Just action ← lookup a actionAndTmkAction = pure (TmkAction action)
    | Just letter ← lookup a actionAndRedirect = letterToTmkLetter' errorL state letter
  letterToTmkLetter' _ _ (Modifiers _ []) = pure (TmkAction "NO")
  letterToTmkLetter' _ _ (Modifiers effect [modifier])
    |  effect ≡ defaultModifierEffect modifier = pure (TmkModifier modifier)
  letterToTmkLetter' _ _ (Redirect [] pos)
    | Just posAction ← lookup pos posAndTmkAction = pure (TmkAction posAction)
  letterToTmkLetter' _ state (Redirect modifiers pos)
    | Just posAction ← lookup pos posAndTmkAction = do
        mods ← intercalate "|" ∘ catMaybes <$> traverse printMod (modifiers ∖ toList state)
        case mods of
          "" → pure (TmkAction posAction)
          _  → pure $ TmkFn ("ACTION_MODS_KEY(" ⊕ mods ⊕ ", KC_" ⊕ posAction ⊕ ")")
    where
      printMod modifier =
        case lookup modifier modifierAndKeycode of
          Just keycode → pure (Just ("MOD_BIT(KC_" ⊕ keycode ⊕ ")"))
          Nothing → Nothing <$ tell [show' modifier ⊕ " is not supported in TMK"]
  letterToTmkLetter' _ state letter
    | action:_ ← mapMaybe posToAction posses = pure (TmkAction action)
    where
      posses = getPosByLetterAndShiftstate letter state defaultFullLayout
      posToAction = flip lookup posAndTmkAction
  letterToTmkLetter' _ _ letter
    | Just tmkLetter ← tmkLetterByMacro [letter] = pure tmkLetter
  letterToTmkLetter' errorL state _ =
      TmkAction "NO" <$ tell [show' errorL ⊕ " is not supported on " ⊕ show' state ⊕ " in TMK"]
  
  tmkLetterByMacro ∷ [Letter] → Maybe TmkLetter
  tmkLetterByMacro letters = TmkMacro mId <$> macro letters <*> pure []
    where
      mId = "LIG_" ⊕ concatMap letterToString letters
      letterToString (Char c)
        | isAscii c ∧ isAlphaNum c = [c]
        | otherwise = "_" ⊕ fromMaybe (printf "U%04X" c) (lookup c charAndString) ⊕ "_"
      letterToString (Unicode c) = printf "_U%04X_" c
      letterToString letter = "_" ⊕ toString letter ⊕ "_"
      macro =
          traverse getPosAndShiftlevel >=>
          groupWith' snd >>>
          (traverse ∘ _1) (traverse (flip lookup modifierAndKeycode) ∘ toList ∘ NE.head ∘ getNonEmpty) >=>
          (traverse ∘ _2 ∘ traverse) (flip lookup posAndTmkAction ∘ fst) >$>
          concatMap (uncurry modsAndCodesToMacro) >>>
          (["SM()", "CM()"] ⧺) ∘ (⧺ ["RM()"])
      getPosAndShiftlevel letter =
          view _keys defaultFullLayout &
          concatMap (\key → map ((,) (view _pos key)) (getValidShiftlevels letter key)) &
          listToMaybe
      getValidShiftlevels letter key =
          filterOnFst (≡ letter) (view _letters key `zip` view _shiftlevels key)
      modsAndCodesToMacro modifiers keycodes =
          map (\m → "D(" ⊕ m ⊕ ")") modifiers ⧺
          map (\c → "T(" ⊕ c ⊕ ")") keycodes ⧺
          map (\m → "U(" ⊕ m ⊕ ")") modifiers
  
  groupLayers ∷ [TmkLayer] → (Int → Int, [(Map Pos TmkLetter, Shiftlevel)])
  groupLayers layers = (flip layerStateAfterGrouping &&& layers') groupedLayers
    where
      groupedLayers = NE.groupWith __tmkLetters layers
      layerStateAfterGrouping initState =
          NE.tail ∘ NE.scanl (\(_, start) gr → (start, start + length gr)) (0, 0) >>>
          map (any (testBit initState) ∘ range ∘ over _2 pred) >>>
          sum ∘ (\xs → zipWith (bool 0 ∘ bit) [0..length xs] xs)
      layers' = map (__tmkLetters ∘ NE.head &&& WithBar ∘ fmap __tmkShiftstate)
  
  addTransLetters ∷ [TmkLayer] → [TmkLayer]
  addTransLetters [] = []
  addTransLetters (x:xs) = addTransLetters' [] x xs
  
  addTransLetters' ∷ [TmkLayer] → TmkLayer → [TmkLayer] → [TmkLayer]
  addTransLetters' low layer high = layer' : nextLayers
    where
      prevLetters = ($ low) $
          filter (`isSubLayer` layer) >>>
          M.unions ∘ map (M.filter (≢ TmkAction "TRNS") ∘ __tmkLetters)
      layer' = over _tmkLetters replaceLetters layer
      replaceLetters = M.unionWith replaceLetter prevLetters
      replaceLetter otherLetter letter
        | otherLetter ≡ letter = TmkAction "TRNS"
        | otherwise = letter
      nextLayers = fromMaybe [] $
          uncurry (addTransLetters' (layer' : low)) <$> uncons high
  
  allTrans ∷ TmkLayer → Bool
  allTrans = all (∈ map TmkAction ["TRNS", "NO"]) ∘ __tmkLetters
  
  isSubLayer ∷ TmkLayer → TmkLayer → Bool
  isSubLayer = isSubState `on` __tmkShiftstate
  
  isSubState ∷ Shiftstate → Shiftstate → Bool
  isSubState = (S.⊆) `on` getSet
  