{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE TypeOperators         #-}

import           Prelude                      hiding (mod)

import           XMonad                       hiding (XConfig (..), get, modify,
                                               put, xmonad)
import           XMonad                       (XConfig (XConfig))
import qualified XMonad                       as X (XConfig (..), xmonad)
import qualified XMonad.StackSet              as W


import           Control.Monad.State

import           Control.Arrow                (first, second)
import           Data.Constraint
import           XMonad.Layout.LayoutModifier (LayoutModifier, ModifiedLayout)


type IsLayout l a = (Read (l a), LayoutClass l a)

class CC m a where
  dict :: forall l. IsLayout l a => m l a -> Dict (IsLayout (m l) a)
  cc :: forall l. m l a -> IsLayout l a :- IsLayout (m l) a
  cc x = Sub $ dict x
  {-# MINIMAL dict #-}

wrapLT :: CC m a
       => (forall l. (LayoutClass l a) => l a -> m l a)
       -> Layout a -> Layout a
wrapLT m (Layout la) = Layout (m la) \\ cc (m la)

newtype (f :. g) l a = O (f (g l) a)
infixr 9 :.

unO :: (f :. g) l a -> f (g l) a
unO (O fgla) = fgla

instance Read (m1 (m2 l) a) => Read ((m1 :. m2) l a) where
  readsPrec i = map (first O) . readsPrec i

instance Show (m1 (m2 l) a) => Show ((m1 :. m2) l a) where
  show = show . unO

instance LayoutClass (m1 (m2 l)) a => LayoutClass ((m1 :. m2) l) a where
  runLayout ws@W.Workspace{W.layout = lay} =
    (second (O <$>) <$>) . runLayout (ws{W.layout = unO lay})
  handleMessage lay =
    fmap (fmap O) . handleMessage (unO lay)
  description = description . unO

instance (CC m1 a, CC m2 a) => CC (m1 :. m2) a where
  dict x = Dict \\ trans (cc (unO x)) (cc undefined)

instance CC Mirror a where dict _ = Dict
instance IsLayout m a => CC (Choose m) a where dict _ = Dict
instance LayoutModifier m a => CC (ModifiedLayout m) a where dict _ = Dict

type Prime = Arr (XConfig Layout)
type Arr a = StateT a IO ()

xmonad :: Prime -> IO ()
xmonad prime = xmonad' =<< execStateT prime (def{X.layoutHook = Layout $ X.layoutHook def})
  where xmonad' :: XConfig Layout -> IO ()
        xmonad' cf@XConfig{ X.layoutHook = Layout l } = X.xmonad cf{ X.layoutHook = l }


modifyLayout :: (CC m Window)
              => (forall l. (LayoutClass l Window) => l Window -> m l Window)
              -> Prime
modifyLayout f = modify $ \c -> c { X.layoutHook = wrapLT f $ X.layoutHook c }

-- imagine there is a function f out there somewhere:
f :: LayoutClass l a => l a -> Mirror (Choose Full l) a
f = Mirror . (Full |||)

main :: IO ()
main = xmonad $ do
  -- this works OK
  modifyLayout (Full ||| )
  modifyLayout Mirror
  -- this doesn't
  modifyLayout $ Mirror . (Full |||)
  -- this does, but is annoying
  modifyLayout $ O . Mirror . (Full |||)
  -- this is the real problem:
  modifyLayout f