-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Steffen Jost <jost@tcs.ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Utils.TH where
-- Common Utility Functions that require TemplateHaskell

-- import Data.Char

import ClassyPrelude.Yesod
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
-- import Control.Monad
-- import Control.Monad.Trans.Class
-- import Control.Monad.Trans.Maybe
-- import Control.Monad.Trans.Except

import Data.List ((!!), foldl)

import Control.Lens
import Control.Monad.Fail

import Utils.I18n

import qualified Data.Char as Char
import Data.Universe (Universe, Finite)
import qualified Data.Map as Map
import qualified Data.Text as Text

import Utils.PathPiece

------------
-- Tuples --
------------

-- Alternatively uses lenses: "^. _3" projects the 3rd component of an n-tuple for any n >=3, requires import Control.Lens
projNI :: Int -> Int -> ExpQ -- generic projection gives I-th element of N-tuple, i.e. snd3 = $(projNI 3 2)  --ghci -fth
-- $(projN n m) :: (t1,..,tn) -> tm (for m<=n)
projNI n i = do
  x <- newName "x"
  let rhs = varE x
  let pat = tupP $ replicate (pred i) wildP ++ varP x : replicate (n-i) wildP
  lam1E pat rhs


-- | Generic projections N-tuples that are actually left-associative pairs
--   i.e. @$(leftAssociativePairProjection c n m :: (..(t1 `c` t2) `c` .. `c` tn) -> tm@ (for m<=n)
leftAssociativePairProjection :: Name -> Int -> Int -> ExpQ
leftAssociativePairProjection constructor n i = do
    x <- newName "x"
    lamE [pat x n] (varE x)
  where
    pat x 1       = varP x
    pat x w
      | w==i      = conP constructor [wildP, varP x]
      | otherwise = conP constructor [pat x (pred w), wildP]


---------------
-- Functions --
---------------

permuteFun :: [Int] -> ExpQ -- generic permutation of function arguments, i.e. $(permuteFun [2,1]) == flip
permuteFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
--         rhs = appE (varE fn) (varE $ xs!!1)
        ln  = length perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..ln] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

altFun :: [Int] -> ExpQ -- generic permutation/repetition of function arguments, i.e. $(permuteFun [2,1]) == flip
altFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
--         rhs = appE (varE fn) (varE $ xs!!1)
        mx  = maximum $ impureNonNull perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..mx] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

-- |
curryN :: Int -> ExpQ
curryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = map varP (fn:xs)
  let tup = tupE (map varE xs)
  let rhs = appE (varE fn) tup
  lamE pat rhs

uncurryN :: Int -> ExpQ
uncurryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = [VarP fn, TupP (map VarP xs)]
  let rhs = foldl AppE (VarE fn) (map VarE xs)
  return $ LamE pat rhs


afterN :: Int -> ExpQ -- apply a function after another of arity N, i.e. $(afterN 1) = (.)
afterN n = do
  f <- newName "f"
  g <- newName "g"
  --let rhs = [|$(curryN n) (g . ($(uncurryN n) f))|]
  lamE [varP g, varP f] [|$(curryN n) $(varE g) . $(uncurryN n) $(varE f)|]


-- Special Show-Instances for Themes
deriveShowWith :: (String -> String) -> Name -> Q [Dec]
deriveShowWith = deriveSimpleWith ''Show 'show

deriveSimpleWith ::  Name -> Name -> (String -> String) -> Name -> Q [Dec]
deriveSimpleWith cls fun strOp ty = do
    (TyConI tyCon)  <- reify ty
    (tyConName, cs) <- case tyCon of
        DataD [] nm [] _ cs _ -> return (nm, cs)
        _ -> fail "deriveShowTheme: tyCon must be a plain datatype enumeration"
    let instanceT = conT cls `appT` conT tyConName
    return <$> instanceD (return []) instanceT [genDecs cs]
  where
    genDecs :: [Con] -> Q Dec
    genDecs cs = funD fun (map genClause cs)

    genClause :: Con -> Q Clause
    genClause (NormalC name []) =
      let pats = [ConP name []]
          body = NormalB $ LitE $ StringL $ strOp $ nameBase name
      in return $ Clause pats body []
    genClause _ = fail "deriveShowTheme: constructors not allowed to have arguments"

embedRenderMessage :: Name -- ^ Foundation type
                   -> Name -- ^ Type to embed into message type
                   -> (Text -> Text) -- ^ Mangle constructor names
                   -> DecsQ
-- ^ @embedRenderMessage ''Foundation ''MessageType mangle@ declares a
--   `RenderMessage Foundation MessageType` instance expecting there
--   to be one constructor for each constructor of @MessageType@ in
--   scope, taking the same arguments:
--
-- > data NewMessage = NewMessageOne | NewMessageTwo | NewMessageThree
-- > data FoundationMessage = MsgOne | MsgThree
-- > data FoundationEvenMessage = MsgTwo
-- >
-- > -- embedRenderMessage ''Foundation ''NewMessage $ drop 2 . splitCamel
-- > instance RenderMessage Foundation NewMessage where
-- >   renderMessage f ls = \case
-- >     NewMessageOne -> renderMessage f ls MsgOne
-- >     NewMessageTwo -> renderMessage f ls MsgTwo
embedRenderMessage f inner mangle = do
  DatatypeInfo{..} <- reifyDatatype inner

  f' <- newName "f"
  ls <- newName "ls"

  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let constr = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP constructorName $ map varP vars) (normalB [e|renderMessage $(varE f') $(varE ls) $(constr)|]) []

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT inner)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ lamCaseE matches) []
      ]
    ]

embedRenderMessageVariant :: Name -- ^ Foundation Type
                          -> Name -- ^ Name of newtype
                          -> (Text -> Text) -- ^ Mangle constructor names
                          -> DecsQ
embedRenderMessageVariant f newT mangle = do
  [ConstructorInfo{ constructorName = newtypeName, constructorFields = [ ConT newtypeInner ] }] <- datatypeCons <$> reifyDatatype newT
  DatatypeInfo{..} <- reifyDatatype newtypeInner

  f' <- newName "f"
  ls <- newName "ls"

  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let body = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP newtypeName [conP constructorName $ map varP vars]) (normalB [e|renderMessage $(varE f') $(varE ls) $(body)|]) []

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT newT)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ lamCaseE matches) []
      ]
    ]


dispatchTH :: Name -- ^ Datatype to pattern match
           -> ExpQ
-- ^ Produces a lambda-case-expression matching all constructors of the named datatype and calling a function (named after the constructor prefixed with @dispatch@) on the fields of each constructor
dispatchTH dType = do
  DatatypeInfo{..} <- reifyDatatype dType
  let
    matches = map mkMatch datatypeCons
    mkMatch ConstructorInfo{..} = do
      pats <- forM constructorFields $ \_ -> newName "x"
      let fName = mkName $ "dispatch" <> nameBase constructorName
      match (conP constructorName $ map varP pats) (normalB $ foldl (\e pat -> e `appE` varE pat) (varE fName) pats) []
  lamCaseE matches


mkI18nWidgetEnum :: String -> FilePath -> DecsQ
mkI18nWidgetEnum (splitCamel -> namebase) basename = do
  itemsAvailable <- i18nWidgetFilesAvailable' basename
  let items = Map.mapWithKey (\k _ -> typPrefix <> unPathPiece k) itemsAvailable
  sequence
    [ dataD (cxt []) dataName [] Nothing
        [ normalC (mkName conName) []
        | (_, conName) <- Map.toAscList items
        ]
        [ derivClause (Just StockStrategy)
            [ conT ''Eq
            , conT ''Ord
            , conT ''Read
            , conT ''Show
            , conT ''Enum
            , conT ''Bounded
            , conT ''Generic
            -- , conT ''Typeable
            ]
        , derivClause (Just AnyclassStrategy)
            [ conT ''Universe
            , conT ''Finite
            , conT ''NFData
            ]
        ]
    , instanceD (cxt []) (conT ''PathPiece `appT` conT dataName)
      [ funD 'toPathPiece
        [ clause [conP (mkName con) []] (normalB . litE . stringL $ repack int) []
        | (int, con) <- Map.toList items
        ]
      , funD 'fromPathPiece
        [ clause [varP $ mkName "t"]
          ( guardedB
            [ (,) <$> normalG [e|$(varE $ mkName "t") == int|] <*> [e|Just $(conE $ mkName con)|]
            | (int, con) <- Map.toList items
            ]) []
        , clause [wildP] (normalB [e|Nothing|]) []
        ]
      ]
    , sigD (mkName $ valPrefix <> "ItemMap") [t|Map Text $(conT dataName)|]
    , funD (mkName $ valPrefix <> "ItemMap")
        [ clause [] (normalB [e| Map.fromList $(listE . map (\(int, con) -> tupE [litE . stringL $ repack int, conE $ mkName con]) $ Map.toList items) |]) []
        ]
    ]
  where
    unPathPiece :: Text -> String
    unPathPiece = repack . mconcat . map (over _head Char.toUpper) . Text.splitOn "-"

    dataName = mkName $ typPrefix <> "Item"

    typPrefix = concat $ over (takingWhile Char.isLower $ _head . traverse) Char.toUpper namebase
    valPrefix = concat $ over (takingWhile Char.isUpper $ _head . traverse) Char.toLower namebase