{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE TemplateHaskell #-}

module Foreign.CPP.HsStruct.HsStdVariant
  ( deriveHsStdVariantUnsafe
  ) where

import Language.Haskell.TH

import Foreign
import Foreign.C.Types
import Foreign.CPP.Addressable
import Foreign.CPP.Marshallable

deriveHsStdVariantUnsafe
  :: String -> Int -> Int -> String -> TypeQ -> Q [Dec]
deriveHsStdVariantUnsafe cppName sizeVal alignmentVal hsName hsType = do
  typeName <- do
    v <- lookupTypeName hsName
    case v of
      Just val -> return val
      Nothing -> fail ("Unable to find type " ++ hsName)
  -- Get the constructors from the data type
  constructors <- do
    info <- reify typeName
    case info of
      -- start zip from `1` since `0` corresponds to `monostate`
      TyConI (DataD _ _ _ _ cons' _) -> return $ zip [1..] cons'
      _ -> fail ("Unsupported type for HsStdVariant " ++ hsName)

  addressableInst <- instanceD
    (cxt [])
    [t| Addressable $hsType |]
    []

  let (cppPeekValName, hsPeekValName) = importNames "peek"
  peekValImport <- forImpD cCall unsafe cppPeekValName hsPeekValName
    [t| Ptr () -> Ptr CInt -> IO (Ptr ()) |]

  let (cppPokeValName, hsPokeValName) = importNames "poke"
  pokeValImport <- forImpD cCall unsafe cppPokeValName hsPokeValName
    [t| Ptr () -> Ptr () -> CInt -> IO () |]

  (peekMatches, pokeClauses) <-
    unzip <$> mapM (getConData hsPokeValName) constructors

  -- Storable
  let
    sizeOfFn = funD (mkName "sizeOf")
      [ clause [wildP] (normalB $ litE $ integerL $ fromIntegral sizeVal) [] ]

    alignmentFn = funD (mkName "alignment")
      [ clause [wildP] (normalB $
          litE $ integerL $ fromIntegral alignmentVal) []
      ]

    pokeFn = funD (mkName "poke") pokeClauses

    idxN = mkName "idx"
    idxpN = mkName "idx_p"
    peekN = mkName "peek"

    -- Peek by calling `hsPeekValName` to extract idx and ptr
    -- Unknown returned index results in `error`
    peekFn = funD peekN
      [clause [varP pN]
        (normalB $ varE (mkName "alloca") `appE` lamE [varP idxpN] (doE
          [ bindS
            (varP valpN)
            (varE hsPeekValName `appE`
              parensE (varE castPtrN `appE` varE pN) `appE`
              varE idxpN)
          , bindS (varP idxN) (varE peekN `appE` varE idxpN)
          , noBindS (caseE (varE idxN) (peekMatches ++
              [ match wildP (normalB $
                varE (mkName "error") `appE` parensE
                  (infixApp
                    (stringE ("Unable to peek " ++ hsName ++ " with index "))
                    (varE $ mkName "++")
                    (varE (mkName "show") `appE` varE idxN)
                  )
                ) []
              ])
            )
          ])) []
      ]

  storableInst <- instanceD
    (cxt [])
    [t| Storable $hsType |]
    [sizeOfFn, alignmentFn, pokeFn, peekFn]


  -- Constructible
  constructibleInst <- instanceD
    (cxt [])
    [t| Constructible $hsType |]
    []

  return
    [ addressableInst
    , constructibleInst
    , peekValImport
    , pokeValImport
    , storableInst
    ]
  where
    pN = mkName "p"
    castPtrN = mkName "castPtr"
    valpN = mkName "__val_p"

    importNames pre = ("std_variant_" ++ base, mkName $ "c_" ++ base)
      where
        base = pre ++ cppName

    getConData hsPokeName (idx, NormalC name [(_, t)] ) = do
        isDestructible <- isInstance (mkName "Destructible") [t]

        let
          -- A match clause for the given index, using the constructor `name`
          -- after `peek`ing the result
          peekMatch = match (litP $ integerL idx) (normalB $ infixApp
              (conE name)
              (varE $ mkName "fmap")
              (varE (mkName "peek") `appE` parensE (varE castPtrN `appE`
                varE valpN))
            )
            []

        let
          valN = mkName "__val"
          withObjFn = if isDestructible then "withCxxObject" else "with"
          -- Matching against the particular constructor `name`,
          -- Using with/withCxxObject to get a pointer to the underlying data
          -- to send down
          pokeClause = clause [varP pN, conP name [varP valN]] (normalB $
            varE (mkName withObjFn) `appE` varE valN `appE` lamE [varP valpN] (
              varE hsPokeName `appE`
              parensE (varE castPtrN `appE` varE pN) `appE`
              parensE (varE castPtrN `appE` varE valpN) `appE`
              litE (integerL idx)
            ))
            []

        return (peekMatch, pokeClause)
    getConData _ _ = fail "Only take unnamed constructors with 1 type"