module Utils.Lens.TH
  ( makeLenses_, makeClassyFor_
  , multifocusG, multifocusL
  ) where

import ClassyPrelude
import Control.Lens
import Control.Lens.Internal.FieldTH
import Language.Haskell.TH

import Numeric.Natural

import Data.Foldable (Foldable(foldl))

-- import Control.Lens.Misc
{-
    NOTE: The code for lensRules_  and makeLenses_ was stolen from package lens-misc-0.0.2.0,
    which was currently unavailable in our stack snapshot.
    See https://github.com/louispan/lens-misc
-}

-- | A 'LensRules' used by 'makeLenses_'.
lensRules_ :: LensRules
lensRules_ = lensRules
    & lensField .~ \_ _ n -> [TopName (mkName ('_':nameBase n))]

-- | Like @lensRules_@, but different class and function name
classyRulesFor_ :: ClassyNamer -> LensRules
classyRulesFor_ clsNamer = classyRules
    & lensClass .~ clsNamer
    & lensField .~ \_ _ n -> [TopName (mkName ('_':nameBase n))]


-- | Build lenses (and traversals) with a sensible default configuration.
-- Works the same as 'makeLenses' except that
-- the resulting lens is also prefixed with an underscore.
--
-- /e.g./
--
-- @
-- data FooBar
--   = Foo { x, y :: 'Int' }
--   | Bar { x :: 'Int' }
-- 'makeLenses' ''FooBar
-- @
--
-- will create
--
-- @
-- _x :: 'Lens'' FooBar 'Int'
-- _x f (Foo a b) = (\\a\' -> Foo a\' b) \<$\> f a
-- _x f (Bar a)   = Bar \<$\> f a
-- _y :: 'Traversal'' FooBar 'Int'
-- _y f (Foo a b) = (\\b\' -> Foo a  b\') \<$\> f b
-- _y _ c\@(Bar _) = pure c
-- @
--
-- @
-- 'makeLenses_' = 'makeLensesWith' 'lensRules_'
-- @
makeLenses_ :: Name -> DecsQ
makeLenses_ = makeFieldOptics lensRules_

-- | like makeClassyFor but only specifies names for class and its function,
--   otherwise lenses are created with underscore like `makeLenses_`
makeClassyFor_ :: Name -> DecsQ
makeClassyFor_ recName = makeFieldOptics (classyRulesFor_ clNamer) recName
    where
      clsName = "Has" <> nameBase recName
      funName = "has" <> nameBase recName
      
      clNamer :: ClassyNamer
      -- clNamer _ = Just (clsName, funName) -- for newer versions >= 4.17
      clNamer _ = Just (mkName clsName, mkName funName)

multifocusG :: Natural -> ExpQ
multifocusG = multifocusOptic
  [e|to . view|]
  (\s a -> [t|Getting $(a) $(s) $(a)|])
  (\s a -> [t|Getter $(s) $(a)|])
  (\doGet _doSet -> [e|to $(doGet)|])

multifocusL :: Natural -> ExpQ
multifocusL = multifocusOptic
  [e|cloneLens|]
  (\s a -> [t|ALens' $(s) $(a)|])
  (\s a -> [t|Lens' $(s) $(a)|])
  (\doGet doSet -> [e|lens $(doGet) $(doSet)|])


multifocusOptic :: _ -> _ -> _ -> _ -> Natural -> ExpQ
multifocusOptic _       _      _     _      0                  = [e|united|]
multifocusOptic doClone _      _     _      1                  = doClone
multifocusOptic doClone alensT lensT lens' (fromIntegral -> n) = do
  ll <- newName "l"
  ls <- replicateM n $ newName "l"
  s  <- newName "s"
  xs <- replicateM n $ newName "x"

  tS  <- newName "s"
  tXs <- replicateM n $ newName "x" :: Q [Name]

  let tup = foldl (\t x -> [t|$(t) $(varT x)|]) (tupleT (length tXs)) tXs
      mkL x = alensT (varT tS) (varT x)

  letE
    [ sigD ll $ foldr (\x t -> [t|$(mkL x) -> $(t)|]) (lensT (varT tS) tup) tXs
    , funD ll
      [ clause
          (map (viewP doClone . varP) ls)
          (normalB $ lens'
             (lamE [varP s] . tupE . flip map ls $ \l -> [e| $(varE s) ^. $(varE l) |])
             (lamE [varP s, tupP $ map varP xs] . foldr (\(x, l) x' -> [e|$(x') & $(varE l) .~ $(varE x)|]) (varE s) $ zip xs ls)
          )
          []
      ]
    ]
    (varE ll)