{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

-- NOTE Because this file will be used as the header of auto-generated macro programs,
--      it can't have any project-specific dependencies (such as `Fix`).
module Axel.Parse.AST where

import Data.IORef (IORef, modifyIORef, newIORef, readIORef)
import Data.Semigroup ((<>))

import System.IO.Unsafe (unsafePerformIO)

handleStringEscapes :: String -> String
handleStringEscapes =
  concatMap $ \case
    '\\' -> "\\\\"
    c -> [c]

-- ******************************
-- ** AST                      **
-- ******************************
-- TODO `Expression` should probably be `Traversable`, use recursion schemes, etc.
--      We should provide `toFix` and `fromFix` functions for macros to take advantage of.
--      (Maybe all macros have the argument automatically `fromFix`-ed to make consumption simpler?)
data Expression ann
  = LiteralChar ann Char
  | LiteralInt ann Int
  | LiteralString ann String
  | SExpression ann [Expression ann]
  | Symbol ann String
  deriving (Eq, Functor, Show)

bottomUpFmap ::
     (Expression ann -> Expression ann) -> Expression ann -> Expression ann
bottomUpFmap f x =
  f $
  case x of
    LiteralChar _ _ -> x
    LiteralInt _ _ -> x
    LiteralString _ _ -> x
    SExpression ann' xs -> SExpression ann' (map (bottomUpFmap f) xs)
    Symbol _ _ -> x

bottomUpTraverse ::
     (Monad m)
  => (Expression ann -> m (Expression ann))
  -> Expression ann
  -> m (Expression ann)
bottomUpTraverse f x =
  f =<<
  case x of
    LiteralChar _ _ -> pure x
    LiteralInt _ _ -> pure x
    LiteralString _ _ -> pure x
    SExpression ann' xs -> SExpression ann' <$> traverse (bottomUpTraverse f) xs
    Symbol _ _ -> pure x

topDownFmap ::
     (Expression ann -> Expression ann) -> Expression ann -> Expression ann
topDownFmap f x =
  case f x of
    LiteralChar _ _ -> x
    LiteralInt _ _ -> x
    LiteralString _ _ -> x
    SExpression ann' xs -> SExpression ann' (map (topDownFmap f) xs)
    Symbol _ _ -> x

-- ******************************
-- ** Sourcemap                **
-- ******************************
data SourcePosition =
  SourcePosition
    { line :: Int
    , column :: Int
    }
  deriving (Eq, Show)

renderSourcePosition :: SourcePosition -> String
renderSourcePosition sourcePosition =
  show (line sourcePosition) <> ":" <> show (column sourcePosition)

data SourceMetadata
  = Nothing
  | FromSource SourcePosition
  deriving (Eq, Show)

-- ******************************
-- ** Internal                 **
-- ******************************
toAxel :: Expression ann -> String
toAxel (LiteralChar _ x) = ['{', x, '}']
toAxel (LiteralInt _ x) = show x
toAxel (LiteralString _ xs) = "\"" <> xs <> "\""
toAxel (SExpression _ (Symbol _ "applyInfix":xs)) =
  "{" <> unwords (map toAxel xs) <> "}"
toAxel (SExpression _ (Symbol _ "list":xs)) =
  "[" <> unwords (map toAxel xs) <> "]"
toAxel (SExpression _ [Symbol _ "quote", x]) = '\'' : toAxel x
toAxel (SExpression _ [Symbol _ "quasiquote", x]) = '`' : toAxel x
toAxel (SExpression _ [Symbol _ "unquote", x]) = '~' : toAxel x
toAxel (SExpression _ [Symbol _ "unquoteSplicing", x]) = "~@" <> toAxel x
toAxel (SExpression _ xs) = "(" <> unwords (map toAxel xs) <> ")"
toAxel (Symbol _ x) = x

-- ******************************
-- ** Quoting                  **
-- ******************************
-- TODO Derive these with Template Haskell (they're currently very brittle)
quoteSourceMetadata :: SourceMetadata -> Expression SourceMetadata
quoteSourceMetadata sm@Nothing = Symbol sm "AST.Nothing"
quoteSourceMetadata sm@(FromSource sourcePosition) =
  SExpression
    sm
    [ Symbol sm "AST.FromSource"
    , SExpression
        sm
        [ Symbol sm "AST.SourcePosition"
        , LiteralInt sm (line sourcePosition)
        , LiteralInt sm (column sourcePosition)
        ]
    ]

quoteParseExpression :: Expression SourceMetadata -> Expression SourceMetadata
quoteParseExpression (LiteralChar ann' x) =
  SExpression
    ann'
    [ Symbol ann' "AST.LiteralChar"
    , quoteSourceMetadata ann'
    , LiteralChar ann' x
    ]
quoteParseExpression (LiteralInt ann' x) =
  SExpression
    ann'
    [Symbol ann' "AST.LiteralInt", quoteSourceMetadata ann', LiteralInt ann' x]
quoteParseExpression (LiteralString ann' x) =
  SExpression
    ann'
    [ Symbol ann' "AST.LiteralString"
    , quoteSourceMetadata ann'
    , LiteralString ann' x
    ]
quoteParseExpression (SExpression ann' xs) =
  SExpression
    ann'
    [ Symbol ann' "AST.SExpression"
    , quoteSourceMetadata ann'
    , SExpression ann' (Symbol ann' "list" : map quoteParseExpression xs)
    ]
quoteParseExpression (Symbol ann' x) =
  SExpression
    ann'
    [ Symbol ann' "AST.Symbol"
    , quoteSourceMetadata ann'
    , LiteralString ann' (handleStringEscapes x)
    ]

-- ******************************
-- ** Macro definition         **
-- ******************************
{-# NOINLINE gensymCounter #-}
gensymCounter :: IORef Int
gensymCounter = unsafePerformIO $ newIORef 0

gensym :: IO (Expression SourceMetadata)
gensym = do
  suffix <- readIORef gensymCounter
  let identifier = "aXEL_AUTOGENERATED_IDENTIFIER_" <> show suffix
  modifyIORef gensymCounter succ
  pure $ Symbol Nothing identifier

-- | This allows splice-unquoting of both `[Expression]`s and `SExpression`s,
--   without requiring special syntax for each.
class ToExpressionList a where
  type Annotation a
  toExpressionList :: a -> [Expression (Annotation a)]

instance ToExpressionList [Expression ann] where
  type Annotation [Expression ann] = ann
  toExpressionList :: [Expression ann] -> [Expression ann]
  toExpressionList = id

-- | Because we do not have a way to statically ensure an `SExpression` is passed
--   (and not another one of `Expression`'s constructors instead),
--   we will error at compile-time if a macro attempts to splice-unquote inappropriately.
instance ToExpressionList (Expression ann) where
  type Annotation (Expression ann) = ann
  toExpressionList :: Expression ann -> [Expression ann]
  toExpressionList (SExpression _ xs) = xs
  toExpressionList x =
    error
      (toAxel x <>
       " cannot be splice-unquoted, because it is not an s-expression!")

wrapCompoundExpressions ::
     [Expression SourceMetadata] -> Expression SourceMetadata
wrapCompoundExpressions stmts =
  SExpression Nothing (Symbol Nothing "begin" : stmts)

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