-- SPDX-FileCopyrightText: 2022 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Test.Util.TH
  ( shouldCompileTo
  , shouldCompileIgnoringInstance
  ) where

import Prelude hiding (Type)

import Data.Generics qualified as SYB
import Language.Haskell.TH (pprint)
import Language.Haskell.TH.Syntax
  (Dec(..), Name, Q, TyVarBndr(..), Type(..), mkName, nameBase, runQ)
import Test.Tasty.HUnit (Assertion, (@?=))
import Text.Show qualified (show)

----------------------------------------------------------------------------
-- TemplateHaskell test helpers
----------------------------------------------------------------------------

shouldCompileTo :: HasCallStack => [Dec] -> Q [Dec] -> Assertion
shouldCompileTo actualDecs expectedQ = do
  expectedDecs <- runQ expectedQ
  PrettyDecs (normalizeDecs actualDecs) @?= PrettyDecs (normalizeDecs expectedDecs)

-- | Same as 'shouldCompileTo', but ignores instance declarations of the given class.
shouldCompileIgnoringInstance :: HasCallStack => Name -> [Dec] -> Q [Dec] -> Assertion
shouldCompileIgnoringInstance className actualDecs expectedQ = do
  expectedDecs <- runQ expectedQ
  let actualDecs' = filter (not . isInstance) actualDecs
  PrettyDecs (normalizeDecs actualDecs') @?= PrettyDecs (normalizeDecs expectedDecs)

  where
    isInstance :: Dec -> Bool
    isInstance = \case
      InstanceD _ _ (ConT t `AppT` _) _ | t == className -> True
      _ -> False

-- | Normalize ASTs to make them comparable.
--
-- By default, quoted ASTs and ASTs with names created using 'newName' will have
-- names with unique IDs.
-- For example:
--
-- > decs <- runQ [d|data D = D { f :: Int } |]
-- > putStrLn $ pprint decs
-- >
-- > -- Will generate this AST:
-- > data D_0 = D_1 { f_2 :: Int }
--
-- To be able to check if two ASTs are equivalent, we have to scrub the unique IDs off all names.
--
-- For convenience, to make the output easier to read, we also erase kind annotations when the kind is '*'.
normalizeDecs :: [Dec] -> [Dec]
normalizeDecs decs =
    SYB.everywhere
      (SYB.mkT fixName . SYB.mkT simplifyType . SYB.mkT simplifyTyVar)
      decs
  where
    fixName :: Name -> Name
    fixName = mkName . nameBase

    simplifyType :: Type -> Type
    simplifyType = \case
      SigT t StarT -> t
      t -> t

    simplifyTyVar :: TyVarBndr -> TyVarBndr
    simplifyTyVar = \case
     KindedTV name StarT -> PlainTV name
     tv -> tv

newtype PrettyDecs = PrettyDecs [Dec]
  deriving newtype Eq

instance Show PrettyDecs where
  show (PrettyDecs decs) = pprint decs