{-# 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 qualified Data.Bits
import Data.Char (isSpace)
import Data.Data (Typeable, cast)
import Data.Foldable (forM_)
import Data.Kind (Type)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (mapMaybe)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import Language.Fiddle.Ast
import Language.Fiddle.Compiler.Backend
import Language.Fiddle.Compiler.Backend.Internal
import Language.Fiddle.Compiler.Backend.Internal.FormattedWriter
import qualified Language.Fiddle.Compiler.Backend.Internal.FragTree as FragTree
import Language.Fiddle.Compiler.Backend.Internal.Writer
import Language.Fiddle.Internal.UnitInterface
import Language.Fiddle.Internal.UnitNumbers
import Language.Fiddle.Types
import Numeric (showHex)
import Options.Applicative
import Text.Printf (printf)

data ImplementationInHeader = ImplementationInHeader

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

-- data StructureInterface = Leaf String String

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 <- gets (Set.member file . 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 :: QualifiedPath String -> Text
qualifiedPathToIdentifier = Text.pack . qualifiedPathToString "__" "_"

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 = qualifiedPathToIdentifier fullPath <> "__get"
        returnType = sizeToType size
        fieldName = basenamePart 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 = qualifiedPathToIdentifier fullPath <> "__set"
        setType = sizeToType size
        fieldName = basenamePart fullPath

    case setType of
      Just rt -> do
        textM $ do
          tell $
            Text.pack $
              printf "static inline void %s(struct %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(struct %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 ::
  Modifier f a ->
  String ->
  QualifiedPath String ->
  N Bits ->
  RegisterBitsTypeRef Checked f a ->
  RegisterBitsDecl Checked f a
pattern DefinedBitsP modifier bitsName bitsFullPath offset typeRef <-
  ( DefinedBits
      { qBitsMetadata =
          Present
            QBitsMetadata
              { bitsSpan =
                  Present
                    FieldSpan
                      { offset = offset
                      },
                bitsFullPath = (basenamePart &&& id -> (bitsName, bitsFullPath))
              },
        definedBitsTypeRef = typeRef,
        definedBitsModifier = (Guaranteed modifier)
      }
    )

writeBitsGet ::
  StructName ->
  String ->
  QualifiedPath String ->
  N Bits ->
  RegisterBitsTypeRef 'Checked I A ->
  M ()
writeBitsGet _ _ _ _ _ = return ()

writeBitsSet ::
  StructName ->
  String ->
  QualifiedPath String ->
  N Bits ->
  RegisterBitsTypeRef 'Checked I A ->
  M ()
writeBitsSet structName bitsName fullPath offset typeRef = do
  text "inline static void "
  text (qualifiedPathToIdentifier fullPath)
  text "__set(\n    struct "
  text structName
  text " *o,\n    "
  typeRefToArgs typeRef
  text ") {\n"

  let shiftArguments =
        zipWith
          (\a b -> if b == 1 then a else a <> " * " <> s b)
          (tail setterArgumentNames)
          (snd $ offsetCoefficients typeRef)
          ++ [s offset | offset /= 0]

  withIndent $ do
    text $ maskValue typeRef

    text "int to_set = value"

    unless (null shiftArguments) $ do
      text " << "
      text $ Text.intercalate " + " shiftArguments

    text ";\n"
    text $ "o->" <> Text.pack bitsName <> " = to_set;\n"
  text "}\n\n"
  where
    offsetCoefficients :: RegisterBitsTypeRef Checked I A -> (Int, [Int])
    offsetCoefficients
      RegisterBitsJustBits
        { justBitsExpr = fromIntegral . trueValue -> sz
        } = (sz, [])
    offsetCoefficients
      RegisterBitsReference
        { bitsRefQualificationMetadata =
            Identity
              ( Present
                  ( ExportedBitsDecl
                      { exportedBitsDeclSizeBits = sz
                      }
                    )
                )
        } = (fromIntegral sz, [])
    offsetCoefficients (RegisterBitsArray tr n _) =
      let (recsz, rest) = offsetCoefficients tr
       in ( recsz * fromIntegral (trueValue n),
            rest ++ [recsz]
          )

    s :: (Show a) => a -> Text
    s = Text.pack . show

    maskValue ref = do
      typeRefToMask ref

    typeRefToMask :: RegisterBitsTypeRef Checked I a -> Text
    typeRefToMask RegisterBitsArray {bitsArrayTypeRef = ref} = typeRefToMask ref
    typeRefToMask
      RegisterBitsJustBits
        { justBitsExpr = fromIntegral . trueValue -> sz
        } = "value &= 0x" <> Text.pack (showHex (((1 :: Int) `Data.Bits.shiftL` (sz :: Int)) - 1) "") <> ";\n"
    typeRefToMask
      RegisterBitsReference
        { bitsRefQualificationMetadata =
            Identity
              ( Present
                  ( ExportedBitsDecl
                      { exportedBitsDeclSizeBits = sz
                      }
                    )
                )
        } =
        let num :: Int
            num = (1 `Data.Bits.shiftL` fromIntegral sz) - 1
         in "value = (typeof(value))(value & 0x" <> Text.pack (showHex num "") <> ");\n"

typeRefToArgs :: RegisterBitsTypeRef 'Checked I A -> M ()
typeRefToArgs reg =
  text
    $ Text.intercalate ",\n    "
    $ zipWith
      (\n t -> t <> " " <> n)
      setterArgumentNames
    $ typeRefToArgs' reg
  where
    typeRefToArgs'
      ( RegisterBitsJustBits
          { justBitsExpr = ConstExpression (LeftV v) _
          }
        ) = [typeForBits v]
    typeRefToArgs'
      ( RegisterBitsReference
          { bitsRefQualificationMetadata = (Identity (Present md))
          }
        ) =
        [ "enum "
            <> qualifiedPathToIdentifier
              (metadataFullyQualifiedPath $ getMetadata md)
        ]
    typeRefToArgs' (RegisterBitsArray tr (ConstExpression (LeftV _) _) _) =
      typeRefToArgs' tr ++ ["int"]

    typeForBits = \case
      64 -> "uint64_t"
      32 -> "uint32_t"
      16 -> "uint16_t"
      8 -> "uint8_t"
      _ -> "unsigned"

-- | Decomposes a type ref into a type name (String) and a list of dimensions
-- (in the case of being an array)
-- decomposeBitsTypeRef :: RegisterBitsTypeRef Checked I A -> (String, [N Unitless])
-- decomposeBitsTypeRef (RegisterBitsJustBits )
writeRegisterBody :: StructName -> QRegMetadata True -> RegisterBody Checked I A -> M ()
writeRegisterBody structName regmeta = walk_ registerWalk
  where
    regName = basenamePart (regFullPath regmeta)

    registerWalk :: forall t. (Walk t, Typeable t) => t I A -> M ()
    registerWalk t = case () of
      ()
        | (Just (DefinedBitsP modifier _ fullPath offset typeRef)) <- castTS t -> do
            sequence_ $
              selectByModifier
                modifier
                ( writeBitsGet structName regName fullPath offset typeRef,
                  writeBitsSet structName regName fullPath offset typeRef
                )

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

getEnumConstants :: [EnumConstantDecl Checked I a] -> [(String, N Unitless)]
getEnumConstants = mapMaybe $ \case
  EnumConstantDecl (identToString -> str) (trueValue -> val) _ -> Just (str, val)
  _ -> Nothing

pattern EnumPattern :: QualifiedPath String -> [(String, N Unitless)] -> FiddleDecl Checked I a
pattern EnumPattern qualifiedPath consts <-
  BitsDecl
    { bitsQualificationMetadata =
        Identity
          ( Present
              ( ExportedBitsDecl
                  { exportedBitsDeclMetadata =
                      Metadata
                        { metadataFullyQualifiedPath = qualifiedPath
                        }
                  }
                )
            ),
      bitsType =
        EnumBitType
          { enumBitBody =
              Identity
                ( EnumBody
                    { enumConsts = getEnumConstants . map undirected -> consts
                    }
                  )
          }
    }

emitEnum :: QualifiedPath String -> [(String, N Unitless)] -> M ()
emitEnum (qualifiedPathToIdentifier -> ident) consts = do
  text $ "enum " <> ident <> " "
  body $ do
    forM_ consts $ \(name, val) -> do
      text $
        ident <> "__" <> Text.pack name <> " = " <> Text.pack (show val) <> ",\n"
  text ";\n\n"

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
  () | Just (EnumPattern path consts) <- castTS t -> do
    checkout sF $ emitEnum path consts
    return Stop
  () | Just (ImportStatement {importPath = path}) <- castTS t -> do
    let header = fst (Text.breakOnEnd "." path) <> "h"
     in checkout hF $ do
          text $ Text.pack $ printf "#include \"%s\"\n" header
          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