module Data.Universe.TH
  ( finiteEnum
  , deriveUniverse
  , deriveFinite
  ) where

import Prelude

import Language.Haskell.TH
import Language.Haskell.TH.Datatype

import Data.Universe
import Data.Universe.Helpers (interleave)

import Control.Monad (unless)

import Data.List (elemIndex, nub)
  
import Control.Lens hiding (universe)
import Data.Generics.Product.Types


-- | Get type var bind name
--
-- Stolen from https://hackage.haskell.org/package/template-haskell-util-0.1.1.0
getTVBName :: TyVarBndr -> Name
getTVBName (PlainTV  name  ) = name
getTVBName (KindedTV name _) = name



finiteEnum :: Name -> DecsQ
-- ^ Declare generic `Enum`- and `Bounded`-Instances given `Finite`- and `Eq`-Instances
finiteEnum tName = do
  DatatypeInfo{..} <- reifyDatatype tName

  let datatype = foldl appT (conT datatypeName) $ map (varT . getTVBName) datatypeVars
      tUniverse = [e|universeF :: [$(datatype)]|]

  [d|
    instance Bounded $(datatype) where
      minBound = head $(tUniverse)
      maxBound = last $(tUniverse)

    instance Enum $(datatype) where
      toEnum n
        | n >= 0
        , n < length $(tUniverse)
        = $(tUniverse) !! n
        | otherwise = error $ "toEnum " ++ $(stringE $ nameBase tName) ++ ": out of bounds"
      fromEnum = fromMaybe (error $ "fromEnum " ++ $(stringE $ nameBase tName) ++ ": invalid `universeF`") . flip elemIndex $(tUniverse)

      enumFrom x = map toEnum [fromEnum x .. fromEnum (maxBound :: $(datatype))]
      enumFromThen x y = map toEnum [fromEnum x, fromEnum y .. fromEnum (maxBound :: $(datatype))]
    |]

deriveUniverse, deriveFinite :: Name -> DecsQ
deriveUniverse tName = view _1 <$> deriveUniverse' [e|interleave|] [e|universe|] ([t|Universe|] `appT`) tName
deriveFinite tName = do
  (decs, iCxt) <- deriveUniverse' [e|concat|] [e|universeF|] ([t|Finite|] `appT`) tName
  fmap concat . sequence $
    [ pure decs
    , do
        DatatypeInfo{..} <- reifyDatatype tName
        pure <$> instanceD (pure iCxt) (appT [t|Finite|] . foldl appT (conT datatypeName) $ map (varT . getTVBName) datatypeVars) []
    ]

deriveUniverse' :: ExpQ -> ExpQ -> (TypeQ -> TypeQ) -> Name -> Q ([Dec], Cxt)
deriveUniverse' interleaveExp universeExp mkCxt tName = do
  DatatypeInfo{..} <- reifyDatatype tName

  let consUniverse ConstructorInfo{..} = do
        unless (null constructorVars) $
          fail "Constructors with variables no supported"

        foldl (\f t -> [e|ap|] `appE` f `appE` sigE universeExp (listT `appT` t)) [e|pure $(conE constructorName)|] $ map pure constructorFields

      typ = foldl (\t bndr -> t `appT` varT (getTVBName bndr)) (conT tName) datatypeVars
      iCxt = map (mkCxt . pure) $ filter (\t -> any (flip (elemOf types) t) usedTVars) fieldTypes
        where usedTVars = filter (\n -> any (`usesVar` n) datatypeCons) $ map getTVBName datatypeVars
              usesVar ConstructorInfo{..} n
                | n `elem` map getTVBName constructorVars = False
                | otherwise = any (elemOf types n) constructorFields
              fieldTypes = nub $ concatMap constructorFields datatypeCons

  iCxt' <- cxt iCxt

  (, iCxt') . pure <$> instanceD (pure iCxt') [t|Universe $(typ)|]
    [ funD 'universe
      [ clause [] (normalB . appE interleaveExp . listE $ map consUniverse datatypeCons) []
      ]
    ]