{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

module Axel.Macros where

import Axel.AST
  ( MacroDefinition
  , Statement(SDataDeclaration, SFunctionDefinition, SMacroDefinition,
          SModuleDeclaration, SNewtypeDeclaration, SPragma, SQualifiedImport,
          SRawStatement, SRestrictedImport, STopLevel, STypeSignature,
          STypeSynonym, STypeclassDefinition, STypeclassInstance,
          SUnrestrictedImport)
  , ToHaskell(toHaskell)
  , functionDefinition
  , name
  )
import Axel.Denormalize (denormalizeStatement)
import qualified Axel.Eff.FileSystem as Effs (FileSystem)
import qualified Axel.Eff.FileSystem as FS
  ( createDirectoryIfMissing
  , withCurrentDirectory
  , withTemporaryDirectory
  , writeFile
  )
import Axel.Eff.Process (StreamSpecification(CreateStreams))
import qualified Axel.Eff.Process as Effs (Process)
import Axel.Eff.Resource (readResource)
import qualified Axel.Eff.Resource as Effs (Resource)
import qualified Axel.Eff.Resource as Res
  ( astDefinition
  , macroDefinitionAndEnvironmentFooter
  , macroDefinitionAndEnvironmentHeader
  , macroScaffold
  )
import Axel.Error (Error(MacroError))
import Axel.Haskell.Language
  ( haskellOperatorSymbols
  , haskellSyntaxSymbols
  , isOperator
  )
import Axel.Haskell.Prettify (prettifyHaskell)
import Axel.Haskell.Stack (interpretFile)
import Axel.Normalize (normalizeStatement)
import qualified Axel.Parse as Parse
  ( Expression(LiteralChar, LiteralInt, LiteralString, SExpression,
           Symbol)
  , parseMultiple
  , programToTopLevelExpressions
  , syntaxSymbols
  , topLevelExpressionsToProgram
  )
import Axel.Utils.Display (Delimiter(Newlines), delimit)
import Axel.Utils.Function (uncurry3)
import Axel.Utils.Recursion (Recursive(bottomUpTraverse), exhaustM)
import Axel.Utils.String (replace)

import Control.Lens.Cons (snoc)
import Control.Lens.Operators ((%~), (^.))
import Control.Lens.Tuple (_1, _2)
import Control.Monad (foldM)
import Control.Monad.Freer (Eff, Members)
import Control.Monad.Freer.Error (throwError)
import qualified Control.Monad.Freer.Error as Effs (Error)

import Data.Function ((&))
import Data.List (foldl')
import Data.List.NonEmpty (NonEmpty, nonEmpty)
import qualified Data.List.NonEmpty as NE (head, map, toList)
import Data.Semigroup ((<>))
import qualified Data.Text as T (isSuffixOf, pack, replace, singleton, unpack)

import System.Exit (ExitCode(ExitFailure))
import System.FilePath ((</>))

hygenisizeMacroName :: String -> String
hygenisizeMacroName oldName =
  let suffix =
        if isOperator oldName
          then "%%%%%%%%%%"
          else "_AXEL_AUTOGENERATED_MACRO_DEFINITION"
      suffixedName =
        if T.pack suffix `T.isSuffixOf` T.pack oldName
          then oldName
          else oldName <> suffix
   in T.unpack $
      foldl'
        (\acc (old, new) -> T.replace (T.singleton old) (T.pack new) acc)
        (T.pack suffixedName)
        (filter
           (\(sym, _) -> sym `notElem` Parse.syntaxSymbols)
           (haskellSyntaxSymbols <> haskellOperatorSymbols))

hygenisizeMacroDefinition :: MacroDefinition -> MacroDefinition
hygenisizeMacroDefinition macroDef =
  macroDef & functionDefinition . name %~ hygenisizeMacroName

generateMacroProgram ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Resource] effs)
  => NonEmpty MacroDefinition
  -> [Statement]
  -> [Parse.Expression]
  -> Eff effs (String, String, String)
generateMacroProgram macroDefs env applicationArgs = do
  astDef <- readResource Res.astDefinition
  scaffold <- getScaffold
  macroDefAndEnv <- (<>) <$> getMacroDefAndEnvHeader <*> getMacroDefAndEnvFooter
  pure (astDef, scaffold, macroDefAndEnv)
  where
    insertDefName =
      let defNamePlaceholder = "%%%MACRO_NAME%%%"
       in replace defNamePlaceholder newMacroName
    oldMacroName = NE.head macroDefs ^. functionDefinition . name
    newMacroName = hygenisizeMacroName oldMacroName
    getMacroDefAndEnvHeader =
      insertDefName <$> readResource Res.macroDefinitionAndEnvironmentHeader
    getMacroDefAndEnvFooter = do
      let hygenicMacroDefs = NE.map hygenisizeMacroDefinition macroDefs
      let source =
            prettifyHaskell $ delimit Newlines $
            map
              toHaskell
              (env <> NE.toList (NE.map SMacroDefinition hygenicMacroDefs))
      footer <-
        insertDefName <$> readResource Res.macroDefinitionAndEnvironmentFooter
      pure $ unlines [source, footer]
    getScaffold =
      let insertApplicationArgs =
            let applicationArgsPlaceholder = "%%%ARGUMENTS%%%"
             in replace applicationArgsPlaceholder (show applicationArgs)
       in prettifyHaskell . insertApplicationArgs . insertDefName <$>
          readResource Res.macroScaffold

expansionPass ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Process, Effs.Resource] effs)
  => Parse.Expression
  -> Eff effs Parse.Expression
expansionPass programExpr =
  Parse.topLevelExpressionsToProgram . map denormalizeStatement <$>
  expandMacros (Parse.programToTopLevelExpressions programExpr)

programToTopLevelExpressions :: Parse.Expression -> [Parse.Expression]
programToTopLevelExpressions (Parse.SExpression (Parse.Symbol "begin":stmts)) =
  stmts
programToTopLevelExpressions _ =
  error "programToTopLevelExpressions must be passed a top-level program!"

topLevelExpressionsToProgram :: [Parse.Expression] -> Parse.Expression
topLevelExpressionsToProgram stmts =
  Parse.SExpression (Parse.Symbol "begin" : stmts)

exhaustivelyExpandMacros ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Process, Effs.Resource] effs)
  => Parse.Expression
  -> Eff effs Parse.Expression
exhaustivelyExpandMacros = exhaustM expansionPass

isStatementNonconflicting :: Statement -> Bool
isStatementNonconflicting (SDataDeclaration _) = True
isStatementNonconflicting (SFunctionDefinition _) = True
isStatementNonconflicting (SPragma _) = True
isStatementNonconflicting (SMacroDefinition _) = True
isStatementNonconflicting (SModuleDeclaration _) = False
isStatementNonconflicting (SNewtypeDeclaration _) = True
isStatementNonconflicting (SQualifiedImport _) = True
isStatementNonconflicting (SRawStatement _) = True
isStatementNonconflicting (SRestrictedImport _) = True
isStatementNonconflicting (STopLevel _) = False
isStatementNonconflicting (STypeclassDefinition _) = True
isStatementNonconflicting (STypeclassInstance _) = True
isStatementNonconflicting (STypeSignature _) = True
isStatementNonconflicting (STypeSynonym _) = True
isStatementNonconflicting (SUnrestrictedImport _) = True

expandMacros ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Process, Effs.Resource] effs)
  => [Parse.Expression]
  -> Eff effs [Statement]
expandMacros topLevelExprs = do
  (stmts, macroDefs) <-
    foldM
      (\acc@(stmts, macroDefs) expr -> do
         expandedExprs <- fullyExpandExpr stmts macroDefs expr
         foldM
           (\acc' expandedExpr -> do
              stmt <- normalizeStatement expandedExpr
              pure $ acc' &
                case stmt of
                  SMacroDefinition macroDef -> _2 %~ flip snoc macroDef
                  _ -> _1 %~ flip snoc stmt)
           acc
           expandedExprs)
      ([], [])
      topLevelExprs
  pure $ stmts <> map (SMacroDefinition . hygenisizeMacroDefinition) macroDefs
  where
    fullyExpandExpr stmts allMacroDefs expr = do
      let program = Parse.topLevelExpressionsToProgram [expr]
      expandedExpr <-
        exhaustM
          (bottomUpTraverse
             (\case
                Parse.SExpression xs ->
                  Parse.SExpression <$>
                  foldM
                    (\acc x ->
                       case x of
                         Parse.SExpression (function:args) ->
                           case lookupMacroDefinitions function allMacroDefs of
                             Just macroDefs ->
                               (acc <>) <$>
                               expandMacroApplication
                                 macroDefs
                                 (filter isStatementNonconflicting stmts)
                                 args
                             Nothing -> pure $ snoc acc x
                         _ -> pure $ snoc acc x)
                    []
                    xs
                x -> pure x))
          program
      pure $ Parse.programToTopLevelExpressions expandedExpr

expandMacroApplication ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Process, Effs.Resource] effs)
  => NonEmpty MacroDefinition
  -> [Statement]
  -> [Parse.Expression]
  -> Eff effs [Parse.Expression]
expandMacroApplication macroDefs auxEnv args = do
  macroProgram <- generateMacroProgram macroDefs auxEnv args
  newSource <- uncurry3 evalMacro macroProgram
  Parse.parseMultiple newSource

lookupMacroDefinitions ::
     Parse.Expression -> [MacroDefinition] -> Maybe (NonEmpty MacroDefinition)
lookupMacroDefinitions identifierExpr =
  nonEmpty . filter (`isMacroBeingCalled` identifierExpr)

isMacroBeingCalled :: MacroDefinition -> Parse.Expression -> Bool
isMacroBeingCalled macroDef identifierExpr =
  case identifierExpr of
    Parse.LiteralChar _ -> False
    Parse.LiteralInt _ -> False
    Parse.LiteralString _ -> False
    Parse.SExpression _ -> False
    Parse.Symbol identifier ->
      macroDef ^. functionDefinition . name == identifier

isMacroDefinitionStatement :: Statement -> Bool
isMacroDefinitionStatement (SMacroDefinition _) = True
isMacroDefinitionStatement _ = False

evalMacro ::
     (Members '[ Effs.Error Error, Effs.FileSystem, Effs.Process] effs)
  => String
  -> String
  -> String
  -> Eff effs String
evalMacro astDefinition scaffold macroDefinitionAndEnvironment =
  FS.withTemporaryDirectory $ \directoryName ->
    FS.withCurrentDirectory directoryName $ do
      let astDirectoryPath = "Axel" </> "Parse"
      let macroDefinitionAndEnvironmentFileName =
            "MacroDefinitionAndEnvironment.hs"
      let scaffoldFileName = "Scaffold.hs"
      FS.createDirectoryIfMissing True astDirectoryPath
      FS.writeFile (astDirectoryPath </> "AST.hs") astDefinition
      FS.writeFile
        macroDefinitionAndEnvironmentFileName
        macroDefinitionAndEnvironment
      FS.writeFile scaffoldFileName scaffold
      interpretFile @'CreateStreams scaffoldFileName "" >>= \case
        (ExitFailure _, _, stderr) ->
          throwError $
          MacroError
            ("Temporary directory: " <> directoryName <> "\n\n" <> "Error:\n" <>
             stderr)
        (_, stdout, _) -> pure stdout