{-# 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
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 qualified Language.Fiddle.Compiler.Backend.Internal.FragTree as FragTree
import Language.Fiddle.Compiler.Backend.Internal.FormattedWriter
import Language.Fiddle.Compiler.Backend.Internal.Writer
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

-- | Header fragment. The top. Starts which include guards and has include
-- statements.
hF :: FileFragment
hF = ("HEADER", FragTree.above FragTree.center)

-- | Structures fragment. The text fragment where structures are defined.
sF :: FileFragment
sF = ("HEADER", FragTree.below (snd hF))

-- | Implementation fragment. This is where function implementations go.
iF :: FileFragment
iF = ("HEADER", FragTree.above (snd fF))

-- | Assert fragment. This is where static asserts go.
aF :: FileFragment
aF = ("HEADER", FragTree.below (snd sF))

-- | Footer fragment. This is wehre the file include endif goes.
fF :: FileFragment
fF = ("HEADER", FragTree.below FragTree.center)

type A = Commented SourceSpan

type I = Identity

type M a = FilesM () FormattedWriter CFileState a

newtype CFileState = CFileState
  { includedFiles :: Set String
  }

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

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
    }

transpile ::
  CBackendFlags ->
  () ->
  FiddleUnit Checked Identity (Commented SourceSpan) ->
  TranspileResult
transpile
  CBackendFlags
    { cHeaderOut = headerFile,
      cSourceOut = sourceFile
    }
  ()
  fiddleUnit = toTranspileResult $ fst $ runFilesM execFormattedWriter () (CFileState mempty) hF run
    where
      toTranspileResult :: Map FilePath Text -> TranspileResult
      toTranspileResult mp =
        TranspileResult $
          Map.mapKeys
            ( \case
                "SOURCE" | Right sourceFile' <- sourceFile -> sourceFile'
                "HEADER" -> headerFile
                k -> k
            )
            mp

      run = do
        checkout hF $
          textM $ do
            tell $ "#ifndef " <> headerGuard <> "\n"
            tell $ "#define " <> headerGuard <> "\n\n"
            tell "#include <stdint.h>\n"

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

        walk (transpileWalk sourceFile headerFile) fiddleUnit ()

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

        checkout 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

pad :: M () -> M ()
pad f = text "\n" *> f <* text "\n"

writeStaticAssert :: Text -> String -> N Bytes -> M ()
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 -> M ()
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 -> M ()
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 -> M ()
writeRegisterBody structName regmeta = walk_ registerWalk
  where
    registerWalk :: forall t. (Walk t, Typeable t) => t I A -> M ()
    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) ->
  M ()

-- | 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 -> M ()
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"

          checkout aF $
            writeStaticAssert structName i off

          checkout 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 -> M () -> M ()
union identifier fn = do
  text "#pragma pack(push, 1)\n"
  text $ "union " <> identifier <> " "
  body fn
  text ";\n"
  text "#pragma pack(pop)\n"

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

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

withIndent :: M a -> M a
withIndent = block incIndent decIndent

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

        checkout 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))
    checkout 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