{-# LANGUAGE OverloadedRecordDot #-}
module Gyehoek.Wasm.Syntax
  ( Module(..)
  , Type
  , Import
  , Tag
  , Global
  , Mem
  , Table
  , Func
  , Data
  , Elem
  , Export
  , Idx
  , TypeIdx
  , FuncIdx
  , GlobalIdx
  , TableIdx
  , MemIdx
  , TagIdx
  , ElemIdx
  , DataIdx
  , LabelIdx
  , LocalIdx
  , FieldIdx
  ) where

import Data.Vector (Vector)
import Data.Word (Word32)
import Data.Binary hiding (putList)
import Data.Bits
import Data.Int (Int8, Int32)
import Data.Binary hiding (putList)
import Data.List (List)
import Data.Generics.Labels
import Control.Lens
import Data.Text (Text)
import Control.Applicative (Alternative(..))
import Data.Foldable (traverse_)
import Control.Monad
import Data.Binary.Put (runPut)
import GHC.Generics (Generic)
import Data.Coerce (coerce)


type Type = RecType
data Import
data Tag
data Global
data Mem
data Table
data Func
data Data
data Elem
data Export
data DataCount
data Code = MkCode { locals :: Vector Local, expr :: Expr }

newtype Idx = MkIdx Word32
  deriving (Show, Generic)

newtype TypeIdx = MkType Idx
  deriving (Show, Generic)
  deriving newtype (Binary)
newtype FuncIdx = MkFunc Idx
  deriving (Show, Generic)
  deriving newtype (Binary)
newtype GlobalIdx = MkGlobal Idx
  deriving (Show, Generic)
newtype TableIdx = MkTable Idx
  deriving (Show, Generic)
newtype MemIdx = MkMem Idx
  deriving (Show, Generic)
newtype TagIdx = MkTag Idx
  deriving (Show, Generic)
newtype ElemIdx = MkElem Idx
  deriving (Show, Generic)
newtype DataIdx = MkData Idx
  deriving (Show, Generic)
newtype LabelIdx = MkLabel Idx
  deriving (Show, Generic)
newtype LocalIdx = MkLocal Idx
  deriving (Show, Generic)
newtype FieldIdx = MkField Idx
  deriving (Show, Generic)

data Module = MkModule
  { types :: Vector Type
  , imports :: Vector Import
  , funcs :: Vector Func
  , tables :: Vector Table
  , mems :: Vector Mem
  , tags :: Vector Tag
  , globals :: Vector Global
  , exports :: Vector Export
  , start :: Maybe FuncIdx
  , elems :: Vector Elem
  , dataCounts :: Vector DataCount
  , code :: Vector Code
  , datas :: Vector Data
  }

instance Semigroup Module where
  m1 <> m2 = MkModule
    { types = m1.types <> m2.types
    , imports = m1.imports <> m2.imports
    , tags = m1.tags <> m2.tags
    , globals = m1.globals <> m2.globals
    , mems = m1.mems <> m2.mems
    , tables = m1.tables <> m2.tables
    , funcs = m1.funcs <> m2.funcs
    , datas = m1.datas <> m2.datas
    , elems = m1.elems <> m2.elems
    , start = m1.start <|> m2.start
    , exports = m1.exports <> m2.exports
    }

instance Monoid Module where
  mempty = MkModule
    { types = mempty
    , imports = mempty
    , funcs = mempty
    , tables = mempty
    , mems = mempty
    , tags = mempty
    , globals = mempty
    , exports = mempty
    , start = Nothing
    , elems = mempty
    , datas = mempty
    }

newtype SectionId = MkSectionId Word8
  deriving newtype (Binary)

pattern SectionCustom    = MkSectionId 0
pattern SectionType      = MkSectionId 1
pattern SectionImport    = MkSectionId 2
pattern SectionFunction  = MkSectionId 3
pattern SectionTable     = MkSectionId 4
pattern SectionMemory    = MkSectionId 5
pattern SectionGlobal    = MkSectionId 6
pattern SectionExport    = MkSectionId 7
pattern SectionStart     = MkSectionId 8
pattern SectionElement   = MkSectionId 9
pattern SectionCode      = MkSectionId 10
pattern SectionData      = MkSectionId 11
pattern SectionDataCount = MkSectionId 12
pattern SectionTag       = MkSectionId 13

newtype RecType = MkRecType { subtypes :: List SubType }

data SubType
  = MkSubType { final :: Bool, supertypes :: List TypeIdx, ct :: CompType }
  deriving (Show, Generic)

data CompType = CompTypeFunc { to :: List ValType, from :: List ValType }
  deriving (Show, Generic)

data ValType
  = ValTypeNum NumType
  deriving (Show, Generic)

data NumType
  = F64 | F32 | I64 | I32
  deriving (Show, Generic)

data Local
data Expr



getULEB128 :: (Integral a, Bits a) => Int -> Get a
getULEB128 bitsBudget = do
    if bitsBudget > 0 then return () else fail "integer representation too long"
    val <- getWord8
    if bitsBudget >= 7 || val .&. 0x7F < 1 `shiftL` bitsBudget then return () else fail "integer too large"
    if not (testBit val 7)
    then return $ fromIntegral val
    else do
        rest <- getULEB128 (bitsBudget - 7)
        return $ (fromIntegral $ val .&. 0x7F) .|. (rest `shiftL` 7)

putULEB128 :: (Integral a, Bits a) => a -> Put
putULEB128 val =
    if val < 128
    then putWord8 $ fromIntegral val
    else do
        putWord8 $ 0x80 + (0x7F .&. fromIntegral val)
        putULEB128 $ val `shiftR` 7

getSLEB128 :: (Integral a, Bits a) => Int -> Get a
getSLEB128 bitsBudget = do
    if bitsBudget > 0 then return () else fail "integer representation too long"
    let toInt8 :: Word8 -> Int8
        toInt8 = fromIntegral
    a <- getWord8
    let mask = (0xFF `shiftL` (bitsBudget - 1)) .&. 0x7F
    if bitsBudget >= 7 || a .&. mask == 0 || a .&. mask == mask then return () else fail "integer too large"
    if not (testBit a 7)
    then return . fromIntegral . toInt8 $ (a .&. 0x7f) .|. ((a .&. 0x40) `shiftL` 1)
    else do
        b <- getSLEB128 (bitsBudget - 7)
        return $ (b `shiftL` 7) .|. (fromIntegral (a .&. 0x7f))

putSLEB128 :: (Integral a, Bits a) => a -> Put
putSLEB128 a = go a
    where
        ext = if a >= 0 then 0 else complement 0
        go x = do
            let 
                r = x `shiftR` 7
                w = x .&. 0x7f
            if r /= ext
            then do
                putWord8 (fromIntegral w .|. 0x80)
                go r
            else
                if (testBit w 6 && a < 0) || (not (testBit w 6) && a >= 0)
                then putWord8 (fromIntegral w)
                else do
                    putWord8 (fromIntegral w .|. 0x80)
                    putWord8 (fromIntegral ext .&. 0x7F)



putSection :: SectionId -> Put -> Put
putSection i contents =
  unless (len == 0) $ do
    put i
    put (lengthOf each contents')
    put contents'
  where
    contents' = runPut contents
    len = lengthOf each contents'

putU32 :: Word32 -> Put
putU32 = putULEB128

putI32 :: Int32 -> Put
putI32 = putSLEB128

putList :: Foldable f => (a -> Put) -> f a -> Put
putList f xs = do
  putU32 (fromIntegral $ length xs)
  traverse_ f xs



putTypeSection :: Vector Type -> Put
putTypeSection = putSection SectionType . putList put

putCodeSection :: Vector Type -> Put
putCodeSection = putSection SectionType . putList put

putStartSection :: Maybe FuncIdx -> Put
putStartSection (Just x) = put x
putStartSection Nothing = pure ()

instance Binary Module where
  get = _
  put m = do
    -- magic
    traverse_ putWord8 [0x00, 0x61, 0x73, 0x6d]
    -- version
    traverse_ putWord8 [0x01, 0x00, 0x00, 0x00]
    -- sections
    putTypeSection m.types
    putStartSection m.start

instance Binary RecType where
  get = _
  put rt = case rt.subtypes of
    [st] -> put st
    sts -> do
      putWord8 0x4e
      putList put sts

instance Binary SubType where
  get = _
  put st = do
    putWord8 $ if st.final then 0x4f else 0x50
    putList put st.supertypes
    put st.ct

instance Binary CompType where
  get = _
  put = \case
    CompTypeFunc s t -> do
      putWord8 0x60
      traverseOf_ both (putList put) (s,t)

instance Binary ValType where
  put = \case
    ValTypeNum nt -> put nt

instance Binary NumType where
  get = _
  put = putWord8 . \case
    F64 -> 0x7c
    F32 -> 0x7d
    I64 -> 0x7e
    I32 -> 0x7f

instance Binary Idx where
  get = _
  put = putU32 . coerce