{- HLINT ignore "Use newtype instead of data" -}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
module Gyehoek.Wasm
  (
    -- * syntax
    Module
  , Idx
  , Expr
    -- ** quasiquoters
  , expr
  , Gyehoek.Sexp.sx
  , Gyehoek.Sexp.sxs
  -- * GenMod effect
  , GenMod
  , runGenMod
  , execGenMod
  , defineFunction
  , defineType
  , defineGlobal
  , emit
  , renderModule
  , wat
  , wats
  , defineFunctions
  , defineTypes
  , defineGlobals
  )
  where

import Language.SexpGrammar
  ( SexpIso(..), (>>>) )
import Language.SexpGrammar qualified as Sexp
import Language.SexpGrammar.Generic
import Data.List (List)
import GHC.Generics (Generic)
import Data.Text (Text)
import Effectful
import Numeric.Natural (Natural)
import Effectful.Dispatch.Dynamic
import Effectful.State.Dynamic
import Control.Lens
import Data.Vector.Strict (Vector)
import qualified Data.Vector.Strict as V
import Language.Sexp.Located
import qualified Gyehoek.Sexp
import GHC.IsList (IsList(..))
import Language.Haskell.TH.Quote (QuasiQuoter)
import Data.Data (Data)
import Gyehoek.Sexp (sx)
import Data.Foldable (traverse_)


newtype Module = MkModule { inner :: Vector Sexp }
  deriving (Show, Generic)
  deriving newtype (Semigroup, Monoid)

newtype Expr = MkExpr { inner :: Vector Instr }
  deriving (Show, Generic, Data, Eq)
  deriving newtype (Semigroup, Monoid)

instance IsList Expr where
  type Item Expr = Instr
  fromList = MkExpr . V.fromList
  toList = V.toList . view #inner

newtype Instr = MkInstr { inner :: Sexp }
  deriving (Show, Generic, Data, Eq)

newtype Idx = MkIdx { inner :: Natural }
  deriving (Generic, Data)
  deriving newtype (Show)


-- GenMod

-- | 'GenModState' is a 'Module' paired with the numbers of functions,
-- types, globals, etc. defined in the module.
data GenModState = MkGenModState
  { mod :: Module
  , funcs :: Natural
  , types :: Natural
  , globals :: Natural
  }
  deriving (Show, Generic)

instance Semigroup GenModState where
  m1 <> m2 = MkGenModState
    { mod = m1.mod <> m2.mod
    , funcs = m1.funcs + m2.funcs
    , types = m1.types + m2.types
    , globals = m1.globals + m2.globals
    }

instance Monoid GenModState where
  mempty = MkGenModState
    { mod = mempty
    , funcs = 0
    , types = 0
    , globals = 0
    }

data GenMod :: Effect where
  DefineFunction :: Sexp -> GenMod m Idx
  DefineType :: Sexp -> GenMod m Idx
  DefineGlobal :: Sexp -> GenMod m Idx
  Emit :: Sexp -> GenMod m ()

type instance DispatchOf GenMod = Dynamic

defineFunction :: GenMod :> es => Sexp -> Eff es Idx
defineFunction = send . DefineFunction

defineFunctions :: GenMod :> es => List Sexp -> Eff es (List Idx)
defineFunctions = traverse (send . DefineFunction)

defineType :: GenMod :> es => Sexp -> Eff es Idx
defineType = send . DefineType

defineTypes :: GenMod :> es => List Sexp -> Eff es (List Idx)
defineTypes = traverse (send . DefineType)

defineGlobal :: GenMod :> es => Sexp -> Eff es Idx
defineGlobal = send . DefineGlobal

defineGlobals :: GenMod :> es => List Sexp -> Eff es (List Idx)
defineGlobals = traverse (send . DefineGlobal)

emit :: GenMod :> es => List Sexp -> Eff es ()
emit = traverse_ (send . Emit)

appendAndIncrement
  :: State GenModState :> es
  => LensLike' ((,) Natural) GenModState Natural
  -> Sexp
  -> Eff es Idx
appendAndIncrement l s =
  state \st -> st
    & #mod . #inner <>~ V.singleton s
    & l <<%~ succ
    & _1 %~ MkIdx

runGenMod :: Eff (GenMod : es) a -> Eff es (a, Module)
runGenMod =
  let run = (mapped . _2 %~ view #mod) . runStateLocal (mempty @GenModState)
  in reinterpret run \cases
    _ (DefineFunction s) -> appendAndIncrement #funcs s
    _ (DefineType s) -> appendAndIncrement #types s
    _ (DefineGlobal s) -> appendAndIncrement #globals s
    _ (Emit s) -> #mod . #inner <>= V.singleton s

execGenMod :: Eff (GenMod : es) a -> Eff es Module
execGenMod = fmap snd . runGenMod

renderModule :: Module -> Text
renderModule (MkModule ss) = Gyehoek.Sexp.format [sx|
  (module ##{ss})
|]


-- SexpIso instances

instance SexpIso Idx where
  sexpIso = with \idx ->
    Sexp.integer >>> Sexp.partialOsi f g
    >>> idx
    where
      f n | n < 0 = Left $ Sexp.unexpected "negative"
                        <> Sexp.expected "natural"
          | otherwise = Right $ fromIntegral n
      g = fromIntegral

instance SexpIso Instr where
  sexpIso = with id

instance Gyehoek.Sexp.SpliceSexp Expr where
  spliceSexp = toListOf $ #inner . each . #inner


-- quasiquoters

expr :: QuasiQuoter
expr = Gyehoek.Sexp.makeSxs
  [||MkExpr . V.fromList . (each . #inner %~ Gyehoek.Sexp.stripLocation)
   . fmap (Gyehoek.Sexp.fromSexp @Instr) ||]

wat :: QuasiQuoter
wat = Gyehoek.Sexp.makeSx [|| id ||]

wats :: QuasiQuoter
wats = Gyehoek.Sexp.makeSxs [|| id ||]