{-# LANGUAGE CPP                   #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

-- | Golden test: TH output
module Test.HsBindgen.Golden.Check.TH (check) where

import Control.Monad.State.Strict (State, get, put, runState)
import Data.Foldable qualified as Foldable
import Data.Generics qualified as SYB
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text
import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Ppr qualified as TH
import Language.Haskell.TH.PprLib qualified as TH
import Language.Haskell.TH.Syntax qualified as TH
import System.FilePath (makeRelative, (<.>), (</>))
import Test.Tasty hiding (after)
import Text.SimplePrettyPrint (Pretty (..))
import Text.SimplePrettyPrint qualified as PP

import Clang.Version

import HsBindgen
import HsBindgen.Backend.Hs.Haddock.Documentation qualified as HsDoc
import HsBindgen.Backend.Hs.Name qualified as Hs
import HsBindgen.Backend.HsModule.Render
import HsBindgen.Guasi
import HsBindgen.Imports
import HsBindgen.TH.Internal

import Test.Common.Util.Tasty
import Test.Common.Util.Tasty.Golden (ActualValue (..))
import Test.HsBindgen.Golden.TestCase
import Test.HsBindgen.Resources

{-------------------------------------------------------------------------------
  Tests
-------------------------------------------------------------------------------}

check :: IO TestResources -> TestCase -> TestTree
check testResources test =
    goldenAnsiDiff "th" fixture $ \report ->
      if ghcAtLeast904 then do
        pkgroot <- getTestPackageRoot testResources
        let artefacts = (,) <$> Dependencies <*> FinalDecls
        -- We do not have access to 'Q', and so have to compute the 'getThDecls'
        -- artefact manually.
        (deps, decls) <-
          runTestHsBindgenSuccess report testResources test artefacts

        let thDecls :: Qu [TH.Dec]
            thDecls = uncurry (getThDecls deps) $ Foldable.fold decls

            (st, thdecs) = runQu thDecls

            -- Here we might have headers outside of our package, but in our
            -- test setup that SHOULD cause an error, as we use bundled stdlib.
            -- And we will cause those on CI, which runs tests on different
            -- systems.
            output :: String
            output = unlines $ concat [
                [    "-- addDependentFile "
                  ++ convertWindows (makeRelative pkgroot fp)
                | fp <- st.dependencyFiles
                ]
              , [ "-- " ++ normalizeQuotes l
                | src <- st.cSources, l <- lines src
                ]
              , [ show $ prettyWithDocumentationMap True st.documentationMap d
                | d <- unqualNames thdecs
                ]
              ]

        return $ ActualValue output
      else
        return $ ActualSkipped "ghc < 9.4"
  where
    fixture :: FilePath
    fixture = test.outputDir </> "th" <.> "txt"

    -- Clang version 19 uses <> for some reason.
    --
    -- If generating Golden tests from an older clang (e.g. 19) version then
    -- this is needed.
    --
    normalizeQuotes :: String -> String
    normalizeQuotes =
      case clangVersion of
        ClangVersionUnknown _ -> id
        ClangVersion v
          -- Put your clang version here
          | v >= (19, 0 ,0) -> go False
          | otherwise       -> id
      where
        go _ [] = []
        go False ('"':xs) = '<' : go True xs
        go True ('"':xs) = '>' : go False xs
        go inBracket (x:xs) = x : go inBracket xs

{-------------------------------------------------------------------------------
  Internal auxiliary: manipulate output
-------------------------------------------------------------------------------}

-- unqualify names, qualified names are noisy *and* GHC.Base names move
unqualNames :: [TH.Dec] -> [TH.Dec]
unqualNames = SYB.everywhere $ SYB.mkT mangleName
  where
    mangleName :: TH.Name -> TH.Name
    mangleName n | n == ''()             = TH.Name (TH.OccName "Unit") TH.NameS
    mangleName (TH.Name occ TH.NameG {}) = TH.Name occ TH.NameS
    mangleName n                         = n

convertWindows :: FilePath -> FilePath
convertWindows = map f where
  f '\\' = '/'
  f c    = c

{-------------------------------------------------------------------------------
  Internal auxiliary: 'Qu' monad
-------------------------------------------------------------------------------}

-- | Deterministic monad with TH.Quote instance
newtype Qu a = Qu (State QuState a)
  deriving newtype (Functor, Applicative, Monad)

data QuState = QuState{
      dependencyFiles  :: [FilePath]
    , uniquenessNumber :: !Integer
    , cSources         :: [String]
    , documentationMap :: Map TH.DocLoc HsDoc.Comment
    }

emptyQuState :: QuState
emptyQuState = QuState [] 0 [] Map.empty

instance TH.Quote Qu where
    newName n = Qu $ do
        q@QuState{ uniquenessNumber = u } <- get
        put $! q { uniquenessNumber = u + 1 }
        return $ TH.Name (TH.OccName n) (TH.NameU u)

instance Guasi Qu where
    addDependentFile fp = Qu $ do
        q@QuState{ dependencyFiles = depfiles } <- get
        put $! q { dependencyFiles = depfiles ++ [fp] }

    addCSource src = Qu $ do
        q@QuState{ cSources = csources } <- get
        put $! q { cSources = csources ++ [src] }

    -- Note: we could mock these better, if we want to test error reporting
    -- Currently (2025-04-15) we only report missing extensions,
    -- so there isn't much to test.
    extsEnabled = return []
    reportError _ = return ()

    putLocalDoc nm s = Qu $ do
        q@QuState{ documentationMap = docMap } <- get
        put $!
          q { documentationMap =
                Map.insert (TH.DeclDoc $ TH.mkName (Text.unpack $ Hs.getName nm)) s docMap
            }

runQu :: Qu a -> (QuState, a)
runQu (Qu m) = case runState m emptyQuState of
    (x, q) -> (q, x)

{-------------------------------------------------------------------------------
  Internal auxiliary: Template Haskell functions
-------------------------------------------------------------------------------}

-- | This function pretty prints 'TH.Dec' with their associated documentation.
prettyWithDocumentationMap ::
     Bool
  -> Map TH.DocLoc HsDoc.Comment
  -> TH.Dec
  -> TH.Doc
prettyWithDocumentationMap isTop docMap dec =
  case dec of
    -- Compatibility with older TH version (removes braces from case expressions)

    TH.ValD p r ds ->
            formatDecDoc docMap dec
      TH.$$ thCompatValD docMap p r ds

    -- Declarations with associated documentation

    TH.DataD cxt name tyvars mkind cons derivs ->
            formatDecDoc docMap dec
      TH.$$ ppTypeDef docMap "data" cxt (Just name) tyvars mkind cons derivs

    TH.NewtypeD cxt name tyvars mkind con derivs ->
            formatDecDoc docMap dec
      TH.$$ ppTypeDef docMap "newtype" cxt (Just name) tyvars mkind [con] derivs

    TH.ClassD cxt name tyvars fundeps decs ->
            formatDecDoc docMap dec
      TH.$$ thCompatClassD docMap name tyvars fundeps cxt decs

    TH.InstanceD overlap cxt typ decs ->
            formatDecDoc docMap dec
      TH.$$ thCompatInstanceD docMap overlap cxt typ decs

    -- The only reason foreign import pretty printer is overwritten here is
    -- for our benefit of making the type signatures easier to read. It is
    -- not what the default TH pretty printer would print.
    TH.ForeignD (TH.ImportF callconv safety impent name typ) ->
            formatDecDoc docMap dec
      TH.$$ TH.hsep [ TH.text "foreign import"
                    , pprCallconv callconv
                    , pprSafety safety
                    , TH.text (show impent)
                    , TH.ppr name
                    , TH.text "::"
                    ]
      TH.$$ TH.nest nestDepth (thCompatPprType typ)

    _ ->    formatDecDoc docMap dec
      TH.$$ TH.ppr_dec isTop dec

-- | Helper function to lookup and format documentation
--
formatDecDoc :: Map TH.DocLoc HsDoc.Comment -> TH.Dec -> TH.Doc
formatDecDoc docMap thDec =
  case getDecDocLoc thDec >>= (`Map.lookup` docMap) of
    Nothing -> TH.empty
    Just c  -> pure
             $ PP.runCtxDoc PP.defaultContext (pretty (TopLevelComment c))

{-------------------------------------------------------------------------------
  TH Compatibility Functions

  These functions provide compatibility with older Template Haskell versions.
  They are extracted here to keep the main logic clean.
-------------------------------------------------------------------------------}

-- | Compatible version of 'TH.ValD' pretty-printing (removes braces from case
-- expressions)
--
thCompatValD ::
     Map TH.DocLoc HsDoc.Comment
  -> TH.Pat
  -> TH.Body
  -> [TH.Dec]
  -> TH.Doc
thCompatValD docMap p r ds =
        TH.ppr p TH.<+> thCompatPprBody docMap True r
  TH.$$ thCompatWhereClause docMap ds

-- | Compatible version of 'TH.ClassD' pretty-printing
--
thCompatClassD ::
     TH.PprFlag a
  => Map TH.DocLoc HsDoc.Comment
  -> TH.Name
  -> [TH.TyVarBndr a]
  -> [TH.FunDep]
  -> TH.Cxt
  -> [TH.Dec]
  -> TH.Doc
thCompatClassD docMap name tyvars fundeps cxt decs =
        TH.text "class" TH.<+> TH.pprCxt cxt
                       TH.<+> TH.ppr name
                       TH.<+> TH.hsep (map TH.ppr tyvars)
                       TH.<+> TH.ppr fundeps
  TH.$$ thCompatWhereClause docMap decs

-- | Compatible version of 'TH.InstanceD' pretty-printing
--
thCompatInstanceD ::
     Map TH.DocLoc HsDoc.Comment
  -> Maybe TH.Overlap
  -> TH.Cxt
  -> TH.Type
  -> [TH.Dec]
  -> TH.Doc
thCompatInstanceD docMap overlap cxt typ decs =
        TH.text "instance" TH.<+> maybe TH.empty TH.ppr_overlap overlap
                          TH.<+> TH.pprCxt cxt
                          TH.<+> TH.ppr typ
  TH.$$ thCompatWhereClause docMap decs

-- | Compatible version of where clause pretty-printing (removes braces from
--   case expressions)
--
thCompatWhereClause :: Map TH.DocLoc HsDoc.Comment -> [TH.Dec] -> TH.Doc
thCompatWhereClause _ [] = TH.empty
thCompatWhereClause docMap ds = TH.nest nestDepth
                              $ TH.text "where"
                         TH.<+> TH.vcat (map (prettyWithDocumentationMap False docMap) ds)

-- | Compatible 'TH.Body' pretty-printing
thCompatPprBody ::
     Map TH.DocLoc HsDoc.Comment
  -> Bool
  -> TH.Body
  -> TH.Doc
thCompatPprBody docMap eq body =
  case body of
    TH.GuardedB xs -> TH.nest nestDepth
                   $ TH.vcat
                   $ map (thCompatPprGuarded docMap eqDoc) xs
    TH.NormalB  e  -> eqDoc TH.<+> thCompatPprExp docMap e
  where eqDoc | eq        = TH.equals
              | otherwise = TH.arrow

-- | Compatible guarded expression pretty-printing
--
thCompatPprGuarded ::
     Map TH.DocLoc HsDoc.Comment
  -> TH.Doc
  -> (TH.Guard, TH.Exp)
  -> TH.Doc
thCompatPprGuarded docMap eqDoc (guard', expr) =
  case guard' of
    TH.NormalG guardExpr -> TH.bar
                     TH.<+> thCompatPprExp docMap guardExpr
                     TH.<+> eqDoc
                     TH.<+> thCompatPprExp docMap expr
    TH.PatG    stmts     ->
            TH.bar TH.<+> TH.vcat (TH.punctuate TH.comma $ map TH.ppr stmts)
      TH.$$ TH.nest nestDepth (eqDoc TH.<+> thCompatPprExp docMap expr)

-- | No precedence 'TH.Exp' printer (compatible version)
--
-- We preserve #13856 pattern match
--
thCompatPprExp :: Map TH.DocLoc HsDoc.Comment -> TH.Exp -> TH.Doc
thCompatPprExp docMap expr =
  case expr of
    TH.LamE [] e -> thCompatPprExp docMap e -- #13856
    TH.LamE ps e -> TH.char '\\' TH.<> TH.hsep (map (TH.pprPat TH.appPrec) ps)
                                TH.<+> TH.text "->" TH.<+> thCompatPprExp docMap e
    TH.CaseE e ms   ->
            TH.text "case" TH.<+> thCompatPprExp docMap e TH.<+> TH.text "of"
      TH.$$ semiSep (thCompatPprMatch docMap) ms
    TH.LamCaseE ms  ->
            TH.text "\\case"
      TH.$$ semiSep (thCompatPprMatch docMap) ms
    _               -> TH.ppr expr
  where
    semiSep f = TH.sep . TH.punctuate TH.semi . map f

-- | Pretty-print Callconv for foreign imports
pprCallconv :: TH.Callconv -> TH.Doc
pprCallconv TH.CCall = TH.text "ccall"
pprCallconv TH.StdCall = TH.text "stdcall"
pprCallconv TH.CApi = TH.text "capi"
pprCallconv TH.Prim = TH.text "prim"
pprCallconv TH.JavaScript = TH.text "javascript"

-- | Pretty-print Safety annotations for foreign imports
pprSafety :: TH.Safety -> TH.Doc
pprSafety TH.Safe = TH.text "safe"
pprSafety TH.Unsafe = TH.text "unsafe"
pprSafety TH.Interruptible = TH.text "interruptible"

-- | Compatible version of type pretty-printing for foreign imports
--
-- Formats function types vertically with proper arrow alignment, matching the
-- style used by the preprocessing renderer in HsModule/Render.hs
thCompatPprType :: TH.Type -> TH.Doc
thCompatPprType = go
  where
    go :: TH.Type -> TH.Doc
    go ty =
      case ty of
        TH.ForallT tvs cxt body ->
          TH.text "forall" TH.<+> TH.hsep (map TH.ppr tvs) TH.<> TH.text "."
            TH.<+> (if null cxt then TH.empty else TH.pprCxt cxt TH.<+> TH.text "=>")
            TH.<+> go body

        TH.AppT (TH.AppT TH.ArrowT arg) res ->
          -- Function type: format vertically with arrow alignment
          TH.ppr arg
            TH.$$ TH.nest (-3) (TH.text "->" TH.<+> go res)

        _ -> TH.ppr ty

-- | Compatible match pretty-printing
thCompatPprMatch :: Map TH.DocLoc HsDoc.Comment -> TH.Match -> TH.Doc
thCompatPprMatch docMap (TH.Match p rhs ds) =
        TH.pprMatchPat p TH.<+> thCompatPprBody docMap False rhs
  TH.$$ thCompatWhereClause docMap ds

{-------------------------------------------------------------------------------
  Type Definition Pretty-Printing
-------------------------------------------------------------------------------}

-- | Pretty-print type definitions with documentation
--
ppTypeDef ::
     TH.PprFlag a
  => Map TH.DocLoc HsDoc.Comment
  -> String
  -> TH.Cxt
  -> Maybe TH.Name
  -> [TH.TyVarBndr a]
  -> Maybe TH.Kind
  -> [TH.Con]
  -> [TH.DerivClause]
  -> TH.Doc
ppTypeDef docMap s cxt mbName tyvars mkind cons derivs =
  TH.sep [        TH.text s
           TH.<+> TH.pprCxt cxt
           TH.<+> case mbName of
             Just name -> TH.pprName' TH.Applied name TH.<+> TH.hsep (map TH.ppr tyvars)
             Nothing   -> TH.hsep (map TH.ppr tyvars)
           TH.<+> ksigDoc TH.<+> maybeWhere
         , TH.nest nestDepth (TH.vcat (pref $ map (ppCon docMap) cons))
         , if null derivs
              then TH.empty
              else TH.nest nestDepth
                 $ TH.vcat $ map TH.ppr_deriv_clause derivs
      ]
  where
    pref :: [TH.Doc] -> [TH.Doc]
    pref xs | isGadtDecl = xs
    pref []              = [] -- No constructors; can't happen in H98
    pref (d:ds)          = (TH.char '=' TH.<+> d):map (TH.bar TH.<+>) ds

    maybeWhere :: TH.Doc
    maybeWhere | isGadtDecl = TH.text "where"
               | otherwise  = TH.empty

    isGadtDecl :: Bool
    isGadtDecl = not (null cons) && all isGadtCon cons
        where isGadtCon TH.GadtC{}         = True
              isGadtCon TH.RecGadtC{}      = True
              isGadtCon (TH.ForallC _ _ x) = isGadtCon x
              isGadtCon  _                 = False

    ksigDoc = case mkind of
                Nothing -> TH.empty
                Just k  -> TH.dcolon TH.<+> TH.ppr k

-- | Pretty print 'TH.Con' with documentation
--
ppCon :: Map TH.DocLoc HsDoc.Comment -> TH.Con -> TH.Doc
ppCon docMap con =
  case con of
    TH.RecC name fields ->
            TH.ppr name TH.<+> ppRecordFields docMap fields
      TH.$$ getConNamesDoc docMap [name]

    TH.ForallC _ _ innerCon ->
            TH.ppr con
      TH.$$ getConNamesDoc docMap (get_cons_names innerCon)

    TH.RecGadtC names fields rtype ->
            TH.commaSepApplied names TH.<+> TH.dcolon
                                     TH.<+> ppRecordFields docMap fields
                                     TH.<+> TH.arrow
                                     TH.<+> TH.ppr rtype
      TH.$$ getConNamesDoc docMap names

    _ ->    TH.ppr con
      TH.$$ getConNamesDoc docMap (get_cons_names con)

-- | Pretty-print record fields with documentation
ppRecordFields ::
     Map TH.DocLoc HsDoc.Comment
  -> [TH.VarBangType]
  -> TH.Doc
ppRecordFields docMap fields =
  TH.braces $ TH.sep $ TH.punctuate TH.comma $ map ppField fields
  where
    ppField (fname, bang, ftype) =
      let fieldBase = TH.pprName' TH.Applied fname TH.<+> TH.dcolon TH.<+>
                      TH.pprBangType (bang, ftype)
      in     fieldBase
       TH.$$ getConNamesDoc docMap [fname]

-- | Aggregate all comments for each name and pretty print
getConNamesDoc :: Map TH.DocLoc HsDoc.Comment -> [TH.Name] -> TH.Doc
getConNamesDoc docMap names =
  case foldMap (\n -> Map.lookup (TH.DeclDoc n) docMap) names of
    Nothing -> TH.empty
    Just c  -> pure
             $ PP.runCtxDoc PP.defaultContext (pretty (PartOfDeclarationComment c))

-- | Get constructor names (defined here as it's not in all TH versions)
get_cons_names :: TH.Con -> [TH.Name]
get_cons_names (TH.NormalC n _)     = [n]
get_cons_names (TH.RecC n _)        = [n]
get_cons_names (TH.InfixC _ n _)    = [n]
get_cons_names (TH.ForallC _ _ c)   = get_cons_names c
get_cons_names (TH.GadtC ns _ _)    = ns
get_cons_names (TH.RecGadtC ns _ _) = ns

{-------------------------------------------------------------------------------
  DocLoc extraction
-------------------------------------------------------------------------------}

-- | Get DocLoc for declarations that can have attached documentation.
getDecDocLoc :: TH.Dec -> Maybe TH.DocLoc
getDecDocLoc (TH.FunD n _)                                        = Just $ TH.DeclDoc n
getDecDocLoc (TH.ValD (TH.VarP n) _ _)                            = Just $ TH.DeclDoc n
getDecDocLoc (TH.DataD _ n _ _ _ _)                               = Just $ TH.DeclDoc n
getDecDocLoc (TH.NewtypeD _ n _ _ _ _)                            = Just $ TH.DeclDoc n
#if MIN_VERSION_template_haskell(2,20,0)
getDecDocLoc (TH.TypeDataD n _ _ _)                               = Just $ TH.DeclDoc n
#endif
getDecDocLoc (TH.TySynD n _ _)                                    = Just $ TH.DeclDoc n
getDecDocLoc (TH.ClassD _ n _ _ _)                                = Just $ TH.DeclDoc n
getDecDocLoc (TH.SigD n _)                                        = Just $ TH.DeclDoc n
getDecDocLoc (TH.ForeignD (TH.ImportF _ _ _ n _))                 = Just $ TH.DeclDoc n
getDecDocLoc (TH.ForeignD (TH.ExportF _ _ n _))                   = Just $ TH.DeclDoc n
#if MIN_VERSION_template_haskell(2,22,0)
getDecDocLoc (TH.InfixD _ _ n)                                    = Just $ TH.DeclDoc n
#else
getDecDocLoc (TH.InfixD _ n)                                      = Just $ TH.DeclDoc n
#endif
getDecDocLoc (TH.DataFamilyD n _ _)                               = Just $ TH.DeclDoc n
getDecDocLoc (TH.OpenTypeFamilyD (TH.TypeFamilyHead n _ _ _))     = Just $ TH.DeclDoc n
getDecDocLoc (TH.ClosedTypeFamilyD (TH.TypeFamilyHead n _ _ _) _) = Just $ TH.DeclDoc n
getDecDocLoc (TH.PatSynD n _ _ _)                                 = Just $ TH.DeclDoc n
getDecDocLoc (TH.PatSynSigD n _)                                  = Just $ TH.DeclDoc n
getDecDocLoc (TH.DefaultSigD n _)                                 = Just $ TH.DeclDoc n
getDecDocLoc (TH.InstanceD _ _ t _)                               = Just $ TH.InstDoc t
getDecDocLoc (TH.DataInstD _ _ t _ _ _)                           = Just $ TH.InstDoc t
getDecDocLoc (TH.NewtypeInstD _ _ t _ _ _)                        = Just $ TH.InstDoc t
getDecDocLoc (TH.TySynInstD (TH.TySynEqn _ t _))                  = Just $ TH.InstDoc t
getDecDocLoc _                                                    = Nothing

{-------------------------------------------------------------------------------
  Configuration
-------------------------------------------------------------------------------}

-- | Same as Template Haskell's
--
nestDepth :: Int
nestDepth = 4

-- | Skip TH for older ghc
ghcAtLeast904 :: Bool
#if __GLASGOW_HASKELL__ >=904
ghcAtLeast904 = True
#else
ghcAtLeast904 = False
#endif