{-# LANGUAGE OverloadedStrings #-}

-- | Shared CPP macro machinery for Haskell source preprocessing.
--
-- Compiler progress tools use this module to set identical CPP preprocessor
-- flags (GHC version macros,
-- @MIN_VERSION_*@ macros, and @-D@\/@-U@ options from @.cabal@ files).
module Aihc.Hackage.Cpp
  ( builtinCppMacros,
    cppMacrosFromOptions,
    minVersionMacroNamesFromDeps,
    injectSyntheticCppMacros,
  )
where

import Data.Map.Strict (Map)
import Data.Map.Strict qualified as M
import Data.Maybe (mapMaybe)
import Data.Set (Set)
import Data.Set qualified as S
import Data.Text (Text)
import Data.Text qualified as T

-- | GHC version macros that every preprocessed file sees.
-- Mirrors what GHC itself defines when invoking @cpp@.
builtinCppMacros :: Map Text Text
builtinCppMacros =
  M.fromList
    [ ("__GLASGOW_HASKELL__", "906"),
      ("__GLASGOW_HASKELL_FULL_VERSION__", "\"9.6.7\""),
      ("__GLASGOW_HASKELL_PATCHLEVEL1__", "7"),
      ("__GLASGOW_HASKELL_PATCHLEVEL2__", "0"),
      ("WORD_SIZE_IN_BITS", "64"),
      ("WORD_SIZE_IN_BITS_FLOAT", "64.0"),
      ("SIZEOF_HSWORD", "8"),
      ("SIZEOF_HSDOUBLE", "8"),
      ("SIZEOF_HSFLOAT", "4")
    ]

-- | Build the macro map for the CPP config from a list of @cpp-options@ strings.
--
-- Starts from 'builtinCppMacros' and applies any @-D@ and @-U@ flags found
-- in the option list.
cppMacrosFromOptions :: [String] -> Map Text Text
cppMacrosFromOptions cppOptions =
  foldl apply builtinCppMacros (mapMaybe parseCppMacroOption cppOptions)
  where
    apply macros option =
      case option of
        CppDefine name value -> M.insert name value macros
        CppUndef name -> M.delete name macros

-- | Derive the set of @MIN_VERSION_<pkg>@ macro names from a list of
-- package dependency names.
minVersionMacroNamesFromDeps :: [Text] -> Set Text
minVersionMacroNamesFromDeps = S.fromList . map toMinVersionMacroName
  where
    toMinVersionMacroName pkg = "MIN_VERSION_" <> T.map sanitizePkgChar pkg
    sanitizePkgChar '-' = '_'
    sanitizePkgChar c = c

-- | Prepend synthetic @#define@ lines for @MIN_VERSION_*@ macros to the
-- source, skipping any names that are already explicitly @-D@\/@-U@'d.
--
-- This lets files that guard code with @#if MIN_VERSION_foo(…)@ compile
-- correctly without requiring Cabal's autogenerated @cabal_macros.h@.
injectSyntheticCppMacros :: [String] -> Set Text -> Text -> Text
injectSyntheticCppMacros cppOptions minVersionMacroNames source =
  let existingFromOptions = cppDefinedOrUndefinedFromOptions cppOptions
      reservedCompilerMinVersionNames = S.fromList ["MIN_VERSION_ghc", "MIN_VERSION_GLASGOW_HASKELL"]
      shouldDefine name = not (name `S.member` existingFromOptions)
      compilerMacroLines =
        [ "#define MIN_VERSION_ghc(major1,major2,minor) 1"
        | shouldDefine "MIN_VERSION_ghc"
        ]
          ++ [ "#define MIN_VERSION_GLASGOW_HASKELL(ma,mi,pl1,pl2) 1"
             | shouldDefine "MIN_VERSION_GLASGOW_HASKELL"
             ]
      dynamicLines =
        [ "#define " <> name <> "(major1,major2,minor) 1"
        | name <- S.toAscList minVersionMacroNames,
          not (name `S.member` reservedCompilerMinVersionNames),
          shouldDefine name
        ]
      header =
        if null (compilerMacroLines ++ dynamicLines)
          then ""
          else T.unlines (compilerMacroLines ++ dynamicLines)
   in if T.null header then source else header <> source

-- ---------------------------------------------------------------------------
-- Internal helpers
-- ---------------------------------------------------------------------------

data CppMacroOption
  = CppDefine Text Text
  | CppUndef Text

parseCppMacroOption :: String -> Maybe CppMacroOption
parseCppMacroOption raw =
  let opt = T.strip (stripWrappingQuotes (T.pack raw))
   in case T.stripPrefix "-D" opt of
        Just rest ->
          case T.breakOn "=" rest of
            (name, "") | validMacroName name -> Just (CppDefine name "1")
            (name, value) | validMacroName name -> Just (CppDefine name (T.drop 1 value))
            _ -> Nothing
        Nothing ->
          case T.stripPrefix "-U" opt of
            Just name | validMacroName name -> Just (CppUndef name)
            _ -> Nothing

validMacroName :: Text -> Bool
validMacroName = not . T.null . T.strip

stripWrappingQuotes :: Text -> Text
stripWrappingQuotes txt =
  if T.length txt >= 2 && T.head txt == '"' && T.last txt == '"'
    then T.dropEnd 1 (T.drop 1 txt)
    else txt

cppDefinedOrUndefinedFromOptions :: [String] -> S.Set Text
cppDefinedOrUndefinedFromOptions =
  foldl addName S.empty . mapMaybe parseCppMacroOption
  where
    addName acc option =
      case option of
        CppDefine name _ -> S.insert name acc
        CppUndef name -> S.insert name acc
