{- HLINT ignore "Use newtype instead of data" -}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeepSubsumption #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE RecordPuns #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE DerivingVia #-}
{-# 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
  , declare
  , renderModule
  , wat
  )
  where

import Language.SexpGrammar
  ( SexpIso(..), list, el, (>>>), rest, sym, symbol, (:-) )
import Language.SexpGrammar qualified as Sexp
import Language.SexpGrammar.Generic
import Data.List (List)
import GHC.Generics (Generic, Generically(..))
import Data.Text (Text)
import Data.String (IsString (fromString))
import Text.Printf
import Effectful
import Numeric.Natural (Natural)
import Effectful.Dispatch.Dynamic
import Effectful.State.Dynamic
import Control.Lens
import Data.Generics.Labels
import Data.Vector (Vector)
import Data.String.Interpolate
import qualified Data.Vector as V
import qualified Data.Text as T
import Effectful.Writer.Dynamic
import Control.Applicative (Alternative((<|>)))
import Control.Category qualified as Cat
import Data.Vector.Lens
import Data.Either (fromLeft, fromRight)
import Language.Sexp.Located
import qualified Gyehoek.Sexp
import GHC.IsList (IsList(..))
import Data.Coerce (coerce)
import qualified Control.Category
import Data.Functor (void)
import Language.Haskell.TH.Quote (QuasiQuoter)
import Data.Data (Data)
import Data.Functor.Foldable (cata)
import Gyehoek.Sexp (sx)
import qualified Language.Sexp as SL


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
  Declare :: Sexp -> GenMod m ()

type instance DispatchOf GenMod = Dynamic

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

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

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

declare :: GenMod :> es => Sexp -> Eff es ()
declare = send . Declare

appendAndIncrement
  :: State GenModState :> es
  => LensLike' ((,) _) 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
    _ (Declare 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 ||]