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

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

import Control.Arrow
import Control.Monad (unless)
import Control.Monad.RWS
import Control.Monad.State (State)
import Control.Monad.Trans.Writer (Writer, execWriter)
import Data.Char (isSpace)
import Data.Data (Typeable, cast)
import Data.Foldable (forM_, toList)
import Data.Kind (Type)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Set (Set)
import qualified Data.Set as Set
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.Internal.UnitNumbers
import Language.Fiddle.Types
import Options.Applicative
import Text.Printf (printf)

data ImplementationInHeader = ImplementationInHeader

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

type StructName = Text

newtype Fragment = Fragment Int
  deriving (Eq, Ord)

-- | Header fragment. The top. Starts which include guards and has include
-- statements.
hF :: Fragment
hF = Fragment 0

-- | Structures fragment. The text fragment where structures are defined.
sF :: Fragment
sF = Fragment 25

-- | Implementation fragment. This is where function implementations go.
iF :: Fragment
iF = Fragment 75

-- | Assert fragment. This is where static asserts go.
aF :: Fragment
aF = Fragment 50

-- | Footer fragment. This is wehre the file include endif goes.
fF :: Fragment
fF = Fragment 100

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

type A = Commented SourceSpan

type I = Identity

data St = St

-- | Current local state information while traversing the tree.
data Fmt = Fmt
  { 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 Fmt a)
  deriving newtype (Functor, Applicative, Monad, MonadState Fmt)

indented :: FileM a -> FileM a
indented fn = do
  textM (modify (\(Fmt id p) -> Fmt (id + 1) p))
  fn <* textM (modify (\(Fmt id p) -> Fmt (id - 1) p))

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

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

newtype FileFragments = FileFragments (Map Fragment (FormattedWriter ()))

instance Semigroup FileFragments where
  (FileFragments m1) <> (FileFragments m2) = FileFragments (Map.unionWith (>>) m1 m2)

instance Monoid FileFragments where
  mempty = FileFragments mempty

newtype CFileState = CFileState
  { includedFiles :: Set String
  }

newtype FileM a = FileM {unFileM :: RWS Fragment FileFragments CFileState a}
  deriving newtype (Functor, Applicative, Monad, MonadWriter FileFragments, MonadReader Fragment, MonadState CFileState)

execFileM :: FileM a -> Text
execFileM fm =
  let (_, FileFragments mp) = execRWS (unFileM fm) hF (CFileState mempty)
   in ( execFormattedWriter
          . sequence_
          . Map.elems
      )
        mp

requireInclude :: String -> FileM ()
requireInclude file = do
  b <- (Set.member file) <$> gets includedFiles
  unless b $ do
    under hF $
      text $
        Text.pack $
          printf "#include <%s>\n" file
    modify $ \s -> s {includedFiles = Set.insert file (includedFiles s)}

-- | Writes text to the current fragment context
text :: Text -> FileM ()
text t = flip tellF_ t =<< ask

-- | Writes text to the current fragment context
textM :: FormattedWriter () -> FileM ()
textM t = flip tellFM_ t =<< ask

-- | Executes a file monad within a different fragment.
under :: Fragment -> FileM () -> FileM ()
under fr = local (const fr)

tellF_ :: Fragment -> Text -> FileM ()
tellF_ fp txt = tell $ FileFragments $ Map.singleton fp (tell txt)

tellFM_ :: Fragment -> FormattedWriter () -> FileM ()
tellFM_ fp txtM = tell $ FileFragments $ Map.singleton fp txtM

newtype Files = Files
  { filepaths :: Map FilePath (FileM ())
  }

withFile :: FilePath -> FileM () -> M ()
withFile fp fn = do
  modify
    ( second $ \(Files {filepaths = fps}) ->
        Files
          { filepaths = Map.alter (Just . (>> fn) . fromMaybe (return ())) 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 = pendingLine s <> 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 execFileM fps

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

          -- Pad out the implementation
          under iF $ text "\n"

        walk (transpileWalk sourceFile headerFile) fiddleUnit ()

        withFile headerFile $ do
          under hF $
            textM $ do
              tell "\nstatic_assert(true); // https://github.com/clangd/clangd/issues/1167\n"

          under fF $
            text headerFinal

      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 :: FileM () -> FileM ()
pad f = text "\n" *> f <* text "\n"

writeStaticAssert :: Text -> String -> N Bytes -> FileM ()
writeStaticAssert structName regname off = do
  requireInclude "stddef.h"
  text $
    Text.pack $
      printf
        "\n_Static_assert(offsetof(%s, %s) == 0x%x, \"Offset wrong\");\n"
        structName
        regname
        off

sizeToType :: N Bytes -> Maybe String
sizeToType = \case
  1 -> Just "uint8_t"
  2 -> Just "uint16_t"
  4 -> Just "uint32_t"
  8 -> Just "uint64_t"
  _ -> Nothing

selectByModifier :: Modifier f an -> (a, a) -> [a]
selectByModifier mod (getter, setter) =
  case mod of
    (ModifierKeyword Rw _) -> [getter, setter]
    (ModifierKeyword Ro _) -> [getter]
    (ModifierKeyword Wo _) -> [setter]
    (ModifierKeyword Pr _) -> []

writeRegGet :: StructName -> QRegMetadata True -> FileM ()
writeRegGet
  structType
  ( QRegMetadata
      { regSpan =
          Present
            FieldSpan
              { size = size
              },
        regFullPath = fullPath
      }
    ) = do
    let fnName = "get_" <> qualifiedPathToIdentifier fullPath
        returnType = sizeToType size
        fieldName = NonEmpty.last fullPath

    case returnType of
      Just rt -> do
        textM $ do
          tell $
            Text.pack $
              printf "static inline %s %s(const %s* o) {\n" rt fnName structType
          tell $ Text.pack $ printf "  return o->%s;\n" fieldName
          tell "}\n\n"
      Nothing ->
        -- Return type is not defined, fallback to byte-by-byte copy.
        textM $ do
          tell $
            Text.pack $
              printf
                "static inline void %s(%s* o, uint8_t out[%d]) {\n"
                fnName
                structType
                size
          forM_ [0 .. size - 1] $ \i ->
            tell $ Text.pack $ printf "  out[%d] = o->%s[%d];\n" i fieldName i
          tell "}\n\n"

writeRegSet :: StructName -> QRegMetadata True -> FileM ()
writeRegSet
  structType
  ( QRegMetadata
      { regSpan =
          Present
            FieldSpan
              { size = size
              },
        regFullPath = fullPath
      }
    ) = do
    let fnName = "set_" <> qualifiedPathToIdentifier fullPath
        setType = sizeToType size
        fieldName = NonEmpty.last fullPath

    case setType of
      Just rt -> do
        textM $ do
          tell $
            Text.pack $
              printf "static inline void %s(%s* o, %s v) {\n" fnName structType rt
          tell $ Text.pack $ printf "  o->%s = v;\n" fieldName
          tell "}\n\n"
      Nothing ->
        -- Return type is not defined, fallback to byte-by-byte copy.
        textM $ do
          tell $
            Text.pack $
              printf
                "static inline void %s(%s* o, const uint8_t in[%d]) {\n"
                fnName
                structType
                size
          forM_ [0 .. size - 1] $ \i ->
            tell $ Text.pack $ printf "  o->%s[%d] = in[%d];\n" fieldName i i
          tell "}\n\n"

pattern DefinedBitsP :: String -> NonEmpty String -> N Bits -> RegisterBitsDecl Checked f a
pattern DefinedBitsP bitsName bitsFullPath offset <-
  ( DefinedBits
      { qBitsMetadata =
          Present
            QBitsMetadata
              { bitsSpan =
                  Present
                    FieldSpan
                      { offset = offset
                      },
                bitsFullPath = (NonEmpty.last &&& id -> (bitsName, bitsFullPath))
              }
      }
    )

writeRegisterBody :: StructName -> QRegMetadata True -> RegisterBody Checked I A -> FileM ()
writeRegisterBody structName regmeta = walk_ registerWalk
  where
    registerWalk :: forall t. (Walk t, Typeable t) => t I A -> FileM ()
    registerWalk t = case () of
      ()
        | (Just (DefinedBitsP bitsName fullPath offset)) <- castTS t ->
            text $
              Text.pack $
                printf
                  "// Emit bits %s (%s) at %d\n"
                  bitsName
                  (qualifiedPathToIdentifier fullPath)
                  offset
      _ -> return ()

    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

writeImplementation ::
  StructName ->
  QRegMetadata True ->
  Modifier f a ->
  Maybe (RegisterBody Checked I A) ->
  FileM ()

-- | Register is just padding, don't emit anything
writeImplementation _ (regIsPadding -> True) _ _ = return ()
writeImplementation structName qMeta mod bod = do
  unless (regIsUnnamed qMeta) $
    sequence_ $
      selectByModifier mod (writeRegGet structName qMeta, writeRegSet structName qMeta)

  mapM_ (writeRegisterBody structName qMeta) bod

structBody :: StructName -> ObjTypeBody Checked I A -> FileM ()
structBody structName (ObjTypeBody _ decls _) = do
  forM_ decls $ \(Directed _ decl _) ->
    case decl of
      RegisterDecl
        { qRegMeta = Present regMetadata,
          regIdent = Guaranteed (identToString -> i),
          regModifier = Guaranteed mod,
          regBody = bod,
          regAnnot = ann
        } -> do
          let (Present (FieldSpan off sz)) = regSpan regMetadata

          textM $ do
            emitDocComments ann
            tell (sizeToField i sz)
            tell ";\n"

          under aF $
            writeStaticAssert structName i off

          under iF $
            writeImplementation structName regMetadata mod bod
      TypeSubStructure
        { subStructureBody = Identity bod,
          subStructureName = mname
        } -> do
          text $
            case objBodyType bod of
              Union {} -> "union "
              Struct {} -> "struct "

          body $ structBody structName bod

          textM $ do
            forM_ mname $ \name ->
              tell (Text.pack $ identToString name)

            tell ";\n"
  where
    sizeToField (Text.pack -> f) = \case
      1 -> "volatile uint8_t " <> f
      2 -> "volatile uint16_t " <> f
      4 -> "volatile uint32_t " <> f
      8 -> "volatile uint64_t " <> f
      n -> "volatile uint8_t " <> f <> "[" <> Text.pack (show n) <> "]"

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

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

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

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

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

    dropIf _ 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 _ headerFile t _ = case () of
  ()
    | Just
        ( ObjTypeDecl
            { objTypeQualificationMetadata = Identity metadata,
              objTypeBody = Identity objTypeBody,
              objTypeAnnot = a
            }
          ) <-
        castTS t -> do
        let structureType = case objBodyType objTypeBody of
              Union {} -> union
              Struct {} -> struct

        withFile headerFile $ do
          under sF $ do
            pad $ do
              textM $ emitDocComments a
              let structName = identifierFor (unwrap metadata)
              structureType structName $ do
                structBody structName objTypeBody
        return Stop
  () | Just (getExportedObjectDecl -> Just e) <- castTS t -> do
    let qname = qualifiedPathToIdentifier (metadataFullyQualifiedPath (getMetadata e))
    withFile headerFile $
      under fF $ do
        text "#define "
        text qname
        text $ Text.pack $ printf " ((%s*)0x%08x)\n" (toLiteralTypeName (exportedObjectDeclType e)) (exportedObjectDeclLocation e)

    return Stop
  _ -> return (Continue ())
  where
    toLiteralTypeName :: ReferencedObjectType -> Text
    toLiteralTypeName (ReferencedObjectType str) = qualifiedPathToIdentifier str
    toLiteralTypeName (ArrayObjectType ro _) = toLiteralTypeName ro

    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

    getExportedObjectDecl :: FiddleDecl Checked I A -> Maybe ExportedObjectDecl
    getExportedObjectDecl (ObjectDecl {objectQualificationMetadata = Identity (Present decl)}) = Just decl
    getExportedObjectDecl _ = Nothing