{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module Language.Fiddle.Compiler.Backend.C (cBackend) where

import Control.Arrow (Arrow (second))
import Control.Exception (TypeError (TypeError))
import Control.Monad.Identity (Identity (Identity))
import Control.Monad.RWS
import Control.Monad.Writer
import Data.Char (isSpace)
import Data.Data (Typeable, cast)
import Data.Foldable (forM_, toList)
import Data.Kind (Type)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe)
import Data.String (IsString)
import Data.Text (Text)
import qualified Data.Text as Text
import Language.Fiddle.Ast
import Language.Fiddle.Compiler.Backend
import Language.Fiddle.Internal.UnitInterface
import Language.Fiddle.Types
import Options.Applicative

data ImplementationInHeader = ImplementationInHeader

data CBackendFlags = CBackendFlags
  { cSourceOut :: Either ImplementationInHeader FilePath,
    cHeaderOut :: FilePath
  }

data FilePosition = FilePosition Int
  deriving (Eq, Ord)

headerPos = FilePosition 0

middlePos = FilePosition 50

footerPos = FilePosition 100

tellLn :: (MonadWriter Text m) => Text -> m ()
tellLn s = tell s >> tell "\n"

type A = Commented SourceSpan

type I = Identity

-- | Current local state information while traversing the tree.
data St = St
  { indentLevel :: Int,
    pendingLine :: Text
  }

newtype M a = M {unM :: RWS () () (St, Files) a}
  deriving newtype (Functor, Applicative, Monad, MonadState (St, Files))

newtype FormattedWriter a = FormattedWriter (RWS () Text St a)
  deriving newtype (Functor, Applicative, Monad, MonadState St)

indented :: FormattedWriter a -> FormattedWriter a
indented fn = do
  modify (\(St id p) -> St (id + 1) p)
  fn <* modify (\(St id p) -> St (id - 1) p)

execFormattedWriter :: FormattedWriter a -> Text
execFormattedWriter ((>> flush) -> FormattedWriter rws) = snd $ execRWS rws () (St 0 "")

flush :: FormattedWriter ()
flush = do
  p <- gets pendingLine
  modify $ \s -> s {pendingLine = ""}
  tell p

data Files = Files
  { filepaths :: Map FilePath (Map FilePosition (FormattedWriter ()))
  }

withFileAt :: FilePath -> FilePosition -> FormattedWriter () -> M ()
withFileAt fp pos wr = do
  modify
    ( second $ \(Files {filepaths = fps}) ->
        Files
          { filepaths =
              Map.alter
                ( \(fromMaybe mempty -> posMap) ->
                    Just $
                      Map.alter (Just . (>> wr) . fromMaybe (return ())) pos posMap
                )
                fp
                fps
          }
    )

instance MonadWriter Text FormattedWriter where
  tell txt = FormattedWriter $ do
    indent <- (`Text.replicate` "  ") <$> gets indentLevel
    let lines = Text.splitOn "\n" txt
    forM_ (init lines) $ \line -> do
      pending <- gets pendingLine
      modify $ \s -> s {pendingLine = ""}
      tell indent
      tell pending
      tell line
      tell "\n"
    modify $ \s -> s {pendingLine = last lines}

  listen (FormattedWriter fn) = FormattedWriter $ listen fn

  pass (FormattedWriter fn) = FormattedWriter $ pass fn

cBackend :: Backend
cBackend =
  Backend
    { backendName = "C",
      backendOptionsParser =
        CBackendFlags
          <$> ( Right
                  <$> strOption
                    ( long "c-source-out"
                        <> short 'o'
                        <> help "Output file for the C source file."
                        <> metavar "OUTPUT"
                    )
                    <|> flag'
                      (Left ImplementationInHeader)
                      ( long "impl-in-header"
                          <> help "Put the whole implementation as static inline functions in the header."
                      )
              )
          <*> strOption
            ( long "c-header-out"
                <> short 'h'
                <> help "Output file for the C header file."
                <> metavar "HEADER_OUT"
            ),
      backendIOMakeState = const $ return (),
      backendTranspile = transpile
    }

toTranspileResult :: Files -> TranspileResult
toTranspileResult Files {filepaths = fps} =
  TranspileResult $
    fmap
      ( execFormattedWriter
          . sequence_
          . Map.elems
      )
      fps

transpile ::
  CBackendFlags ->
  () ->
  FiddleUnit Checked Identity (Commented SourceSpan) ->
  TranspileResult
transpile
  CBackendFlags
    { cHeaderOut = headerFile,
      cSourceOut = sourceFile
    }
  ()
  fiddleUnit = toTranspileResult . snd . fst $ execRWS (unM run) () (St 0 "", Files mempty)
    where
      run :: M ()
      run = do
        withFileAt headerFile headerPos $ do
          tell $ "#ifndef " <> headerGuard <> "\n"
          tell $ "#define " <> headerGuard <> "\n"

        walk (transpileWalk sourceFile headerFile) fiddleUnit ()
        withFileAt headerFile footerPos $ tell headerFinal

        withFileAt headerFile headerPos $ do
          tell "\n#include <stdint.h>\n"

          tell "\nstatic_assert(true); // https://github.com/clangd/clangd/issues/1167\n"

      headerFinal = "\n#endif /* " <> headerGuard <> " */\n"

      headerGuard =
        Text.toUpper $
          Text.replace "." "_" $
            Text.replace "/" "_" $
              Text.pack headerFile

class IsText t where
  toText :: t -> Text

instance IsText String where
  toText = Text.pack

instance IsText Text where
  toText = id

qualifiedPathToIdentifier :: (Foldable f, IsText t) => f t -> Text
qualifiedPathToIdentifier = Text.intercalate "_" . map toText . toList

ensureNL :: FormattedWriter ()
ensureNL = do
  p <- gets pendingLine
  if Text.null p
    then return ()
    else do
      modify $ \s -> s {pendingLine = ""}
      tell p
      tell "\n"

pad :: (IsString t, MonadWriter t m) => m a -> m a
pad f = tell "\n" *> f <* tell "\n"

structBody :: ObjTypeBody Checked I A -> FormattedWriter ()
structBody _ = return ()

struct :: Text -> FormattedWriter () -> FormattedWriter ()
struct identifier fn = do
  tell "#pragma pack(push, 1)\n"
  tell $ "struct " <> identifier <> " "
  body fn
  tell ";\n"
  tell "#pragma pack(pop)\n"

body :: FormattedWriter a -> FormattedWriter a
body f = tell "{\n" *> indented f <* (ensureNL >> tell "}")

identifierFor :: (ExportableDecl d) => d -> Text
identifierFor = qualifiedPathToIdentifier . metadataFullyQualifiedPath . getMetadata

emitDocComments :: (MonadWriter Text m) => A -> m ()
emitDocComments (Commented comments _) =
  mapM_ (\t -> tellLn $ "// " <> t) $
    mapMaybe
      ( \case
          (DocComment t) -> Just (trimDocComment t)
          _ -> Nothing
      )
      comments
  where
    trimDocComment =
      Text.dropWhileEnd isSpace
        . Text.dropWhile isSpace
        . dropIf (== '*')
        . Text.dropWhile isSpace

    dropIf fn t | Text.null t = mempty
    dropIf fn t =
      if fn (Text.head t)
        then Text.tail t
        else t

transpileWalk :: Either ImplementationInHeader FilePath -> FilePath -> (forall t. (Walk t, Typeable t) => t I A -> () -> M (WalkContinuation ()))
transpileWalk sourceFile headerFile t _ = case () of
  ()
    | Just
        ( ObjTypeDecl
            { objTypeQualificationMetadata = Identity metadata,
              objTypeIdent = (identToString -> identifier),
              objTypeBody = Identity objTypeBody,
              objTypeAnnot = a
            }
          ) <-
        castTS t -> do
        withFileAt headerFile middlePos $ do
          pad $ do
            emitDocComments a
            struct (identifierFor (unwrap metadata)) $ do
              structBody objTypeBody
        return Stop
  _ -> return (Continue ())
  where
    castTS ::
      forall (t' :: SynTree) (t :: StagedSynTree) (f :: Type -> Type) (a :: Type).
      ( Typeable t',
        Typeable t,
        Typeable f,
        Typeable a
      ) =>
      t' f a ->
      Maybe (t Checked f a)
    castTS = cast