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)


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 pure 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 = deriveUniverse' [e|interleave|] [e|universe|]
deriveFinite tName = fmap concat . sequence $
  [ deriveUniverse' [e|concat|] [e|universeF|] tName
  , do
      DatatypeInfo{..} <- reifyDatatype tName
      [d|instance Finite $(foldl appT (conT datatypeName) $ map pure datatypeVars)|]
  ]

deriveUniverse' :: ExpQ -> ExpQ -> Name -> DecsQ
deriveUniverse' interleaveExp universeExp tName = do
  DatatypeInfo{..} <- reifyDatatype tName

  let datatype = foldl appT (conT datatypeName) $ map pure datatypeVars
      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

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