{-# LANGUAGE UndecidableInstances #-} -- for various stuff
{-# LANGUAGE AllowAmbiguousTypes #-} -- for type-level sum type handling

module Binrep.Put where

import Binrep.BLen ( BLen(blen) )
import Binrep.CBLen ( IsCBLen(CBLen), cblen )
import Bytezap.Poke
import Raehik.Compat.Data.Primitive.Types ( Prim', sizeOf )
import Binrep.Util.ByteOrder
import Raehik.Compat.Data.Primitive.Types.Endian ( ByteSwap )
import Binrep.Common.Via.Prim ( ViaPrim(..) )

import Data.ByteString qualified as B

import Binrep.Common.Class.TypeErrors ( ENoSum, ENoEmpty )
import GHC.TypeLits ( TypeError, KnownNat )

import GHC.Generics
import Generic.Data.Function.FoldMap
import Generic.Data.MetaParse.Cstr ( Raw, ParseCstrTo )
import Generic.Type.Assert

import Control.Monad.ST ( RealWorld )

import Binrep.Put.Struct ( PutC(putC) )

import Rerefined.Refine
import Rerefined.Predicate.Logical.And

import Data.Word
import Data.Int
import Data.Void
import Binrep.Common.Via.Generically.NonSum

type Putter = Poke RealWorld
class Put a where put :: a -> Putter

runPut :: (BLen a, Put a) => a -> B.ByteString
runPut :: forall a. (BLen a, Put a) => a -> ByteString
runPut a
a = Int -> Poke RealWorld -> ByteString
unsafeRunPokeBS (a -> Int
forall a. BLen a => a -> Int
blen a
a) (a -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put a
a)

-- | Serialize generically using generic 'foldMap'.
instance GenericFoldMap Put where
    type GenericFoldMapM Put = Putter
    type GenericFoldMapC Put a = Put a
    genericFoldMapF :: forall a. GenericFoldMapC Put a => a -> GenericFoldMapM Put
genericFoldMapF = a -> Poke RealWorld
a -> GenericFoldMapM Put
forall a. Put a => a -> Poke RealWorld
put

-- | Serialize a term of the non-sum type @a@ via its 'Generic' instance.
putGenericNonSum
    :: forall a
    .  ( Generic a, GFoldMapNonSum Put (Rep a)
       , GAssertNotVoid a, GAssertNotSum a
    ) => a -> Putter
putGenericNonSum :: forall a.
(Generic a, GFoldMapNonSum Put (Rep a), GAssertNotVoid a,
 GAssertNotSum a) =>
a -> Poke RealWorld
putGenericNonSum = forall {k} (tag :: k) a.
(Generic a, GFoldMapNonSum tag (Rep a)) =>
a -> GenericFoldMapM tag
forall (tag :: Type -> Constraint) a.
(Generic a, GFoldMapNonSum tag (Rep a)) =>
a -> GenericFoldMapM tag
genericFoldMapNonSum @Put

instance
  ( Generic a, GFoldMapNonSum Put (Rep a)
  , GAssertNotVoid a, GAssertNotSum a
  ) => Put (GenericallyNonSum a) where
    put :: GenericallyNonSum a -> Poke RealWorld
put = a -> Poke RealWorld
forall a.
(Generic a, GFoldMapNonSum Put (Rep a), GAssertNotVoid a,
 GAssertNotSum a) =>
a -> Poke RealWorld
putGenericNonSum (a -> Poke RealWorld)
-> (GenericallyNonSum a -> a)
-> GenericallyNonSum a
-> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenericallyNonSum a -> a
forall a. GenericallyNonSum a -> a
unGenericallyNonSum

-- | Serialize a term of the sum type @a@ via its 'Generic' instance.
putGenericSum
    :: forall sumtag a
    .  ( Generic a, GFoldMapSum Put sumtag (Rep a)
       , GAssertNotVoid a, GAssertSum a
    ) => ParseCstrTo sumtag Putter -> a -> Putter
putGenericSum :: forall {k} (sumtag :: k) a.
(Generic a, GFoldMapSum Put sumtag (Rep a), GAssertNotVoid a,
 GAssertSum a) =>
ParseCstrTo sumtag (Poke RealWorld) -> a -> Poke RealWorld
putGenericSum = forall {k1} {k2} (tag :: k1) (sumtag :: k2) a.
(Generic a, GFoldMapSum tag sumtag (Rep a)) =>
ParseCstrTo sumtag (GenericFoldMapM tag)
-> a -> GenericFoldMapM tag
forall (tag :: Type -> Constraint) (sumtag :: k) a.
(Generic a, GFoldMapSum tag sumtag (Rep a)) =>
ParseCstrTo sumtag (GenericFoldMapM tag)
-> a -> GenericFoldMapM tag
genericFoldMapSum @Put @sumtag

-- | Serialize a term of the sum type @a@ via its 'Generic' instance, without
--   pre-parsing constructor names.
putGenericSumRaw
    :: forall a
    .  ( Generic a, GFoldMapSum Put Raw (Rep a)
       , GAssertNotVoid a, GAssertSum a
    ) => (String -> Putter) -> a -> Putter
putGenericSumRaw :: forall a.
(Generic a, GFoldMapSum Put Raw (Rep a), GAssertNotVoid a,
 GAssertSum a) =>
(String -> Poke RealWorld) -> a -> Poke RealWorld
putGenericSumRaw = forall {k} (tag :: k) a.
(Generic a, GFoldMapSum tag Raw (Rep a)) =>
(String -> GenericFoldMapM tag) -> a -> GenericFoldMapM tag
forall (tag :: Type -> Constraint) a.
(Generic a, GFoldMapSum tag Raw (Rep a)) =>
(String -> GenericFoldMapM tag) -> a -> GenericFoldMapM tag
genericFoldMapSumRaw @Put

newtype ViaPutC a = ViaPutC { forall a. ViaPutC a -> a
unViaPutC :: a }
instance (PutC a, KnownNat (CBLen a)) => Put (ViaPutC a) where
    {-# INLINE put #-}
    put :: ViaPutC a -> Poke RealWorld
put = Int -> Poke RealWorld -> Poke RealWorld
forall s. Int -> Poke s -> Poke s
fromStructPoke (forall a. KnownNat (CBLen a) => Int
forall {k} (a :: k). KnownNat (CBLen a) => Int
cblen @a) (Poke RealWorld -> Poke RealWorld)
-> (ViaPutC a -> Poke RealWorld) -> ViaPutC a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Poke RealWorld
forall a. PutC a => a -> Poke RealWorld
putC (a -> Poke RealWorld)
-> (ViaPutC a -> a) -> ViaPutC a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ViaPutC a -> a
forall a. ViaPutC a -> a
unViaPutC

-- use ViaPutC over this, but should be semantically identical
instance Prim' a => Put (ViaPrim a) where
    put :: ViaPrim a -> Poke RealWorld
put = Int -> Poke RealWorld -> Poke RealWorld
forall s. Int -> Poke s -> Poke s
fromStructPoke (a -> Int
forall a. Prim a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) (Poke RealWorld -> Poke RealWorld)
-> (ViaPrim a -> Poke RealWorld) -> ViaPrim a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ViaPrim a -> Poke RealWorld
forall a. PutC a => a -> Poke RealWorld
putC
    {-# INLINE put #-}

instance TypeError ENoEmpty => Put Void where put :: Void -> Poke RealWorld
put = Void -> Poke RealWorld
forall a. HasCallStack => a
undefined
instance TypeError ENoSum => Put (Either a b) where put :: Either a b -> Poke RealWorld
put = Either a b -> Poke RealWorld
forall a. HasCallStack => a
undefined

instance Put Putter where put :: Poke RealWorld -> Poke RealWorld
put = Poke RealWorld -> Poke RealWorld
forall a. a -> a
id

-- | Unit type serializes to nothing. How zen.
instance Put () where
    {-# INLINE put #-}
    put :: () -> Poke RealWorld
put = () -> Poke RealWorld
forall a. Monoid a => a
mempty

instance (Put l, Put r) => Put (l, r) where
    {-# INLINE put #-}
    put :: (l, r) -> Poke RealWorld
put (l
l, r
r) = l -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put l
l Poke RealWorld -> Poke RealWorld -> Poke RealWorld
forall a. Semigroup a => a -> a -> a
<> r -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put r
r

instance Put a => Put [a] where
    {-# INLINE put #-}
    put :: [a] -> Poke RealWorld
put = [Poke RealWorld] -> Poke RealWorld
forall a. Monoid a => [a] -> a
mconcat ([Poke RealWorld] -> Poke RealWorld)
-> ([a] -> [Poke RealWorld]) -> [a] -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Poke RealWorld) -> [a] -> [Poke RealWorld]
forall a b. (a -> b) -> [a] -> [b]
map a -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put

instance Put B.ByteString where
    {-# INLINE put #-}
    put :: ByteString -> Poke RealWorld
put = ByteString -> Poke RealWorld
byteString

-- | 8-bit (1-byte) words do not require byte order in order to precisely
--   define their representation.
deriving via ViaPutC Word8 instance Put Word8

-- | 8-bit (1-byte) words do not require byte order in order to precisely
--   define their representation.
deriving via ViaPutC  Int8 instance Put  Int8

-- | Byte order is irrelevant for 8-bit (1-byte) words.
deriving via Word8 instance Put (ByteOrdered end Word8)

-- | Byte order is irrelevant for 8-bit (1-byte) words.
deriving via  Int8 instance Put (ByteOrdered end  Int8)

-- ByteSwap is required on opposite endian platforms, but we're not checking
-- here, so make sure to keep it on both.
-- Stick with ViaPrim here because ByteOrdered is connected to it.
deriving via ViaPrim (ByteOrdered LittleEndian a)
    instance (Prim' a, ByteSwap a) => Put (ByteOrdered LittleEndian a)
deriving via ViaPrim (ByteOrdered    BigEndian a)
    instance (Prim' a, ByteSwap a) => Put (ByteOrdered    BigEndian a)

-- | Put types refined with multiple predicates by wrapping the left
--   predicate with the right. LOL REALLY?
instance Put (Refined pr (Refined pl a)) => Put (Refined (pl `And` pr) a) where
    put :: Refined (And pl pr) a -> Poke RealWorld
put = Refined pr (Refined pl a) -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put (Refined pr (Refined pl a) -> Poke RealWorld)
-> (Refined (And pl pr) a -> Refined pr (Refined pl a))
-> Refined (And pl pr) a
-> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (p :: k1). a -> Refined p a
forall {k} a (p :: k). a -> Refined p a
unsafeRefine @_ @pr (Refined pl a -> Refined pr (Refined pl a))
-> (Refined (And pl pr) a -> Refined pl a)
-> Refined (And pl pr) a
-> Refined pr (Refined pl a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (p :: k). a -> Refined p a
forall {k} a (p :: k). a -> Refined p a
unsafeRefine @_ @pl (a -> Refined pl a)
-> (Refined (And pl pr) a -> a)
-> Refined (And pl pr) a
-> Refined pl a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Refined (And pl pr) a -> a
forall {k} (p :: k) a. Refined p a -> a
unrefine