{-# LANGUAGE UndecidableInstances #-} -- required for type-level stuff
{-# LANGUAGE OverloadedStrings #-} -- required for refined errors

module Binrep.Type.Prefix.Count where

import Binrep.Util.Prefix
import Binrep
import Control.Monad.Combinators qualified as Monad

import GHC.TypeNats
import Util.TypeNats ( natValInt )
import Refined hiding ( Weaken(..), strengthen )
import Refined.Unsafe ( reallyUnsafeRefine1 )

import Data.Typeable ( Typeable, typeRep )
import Data.Kind

import Data.Foldable qualified as Foldable

data CountPrefix (pfx :: Type)
type CountPrefixed pfx = Refined1 (CountPrefix pfx)

instance (KnownNat (Max pfx), Foldable f, Typeable pfx)
  => Predicate1 (CountPrefix pfx) f where
    validate1 :: forall a. Proxy (CountPrefix pfx) -> f a -> Maybe RefineException
validate1 Proxy (CountPrefix pfx)
p f a
fa
      | f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= forall (n :: Nat). KnownNat n => Int
natValInt @(Max pfx) = Maybe RefineException
success
      | Bool
otherwise = TypeRep -> Text -> Maybe RefineException
throwRefineOtherException (Proxy (CountPrefix pfx) -> TypeRep
forall {k} (proxy :: k -> Type) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Proxy (CountPrefix pfx)
p) Text
"TODO bad"

instance (KnownNat (Max pfx), Foldable f, Typeable pfx)
  => Predicate (CountPrefix pfx) (f a) where
    validate :: Proxy (CountPrefix pfx) -> f a -> Maybe RefineException
validate = Proxy (CountPrefix pfx) -> f a -> Maybe RefineException
forall a. Proxy (CountPrefix pfx) -> f a -> Maybe RefineException
forall {k} k1 (p :: k1) (f :: k -> Type) (a :: k).
Predicate1 p f =>
Proxy p -> f a -> Maybe RefineException
validate1

-- | We can know byte length at compile time _if_ we know it for the prefix and
--   the list-like.
--
-- This is extremely unlikely, because then what counting are we even
-- performing for the list-like? But it's a valid instance.
instance IsCBLen (CountPrefixed pfx f a) where
    type CBLen (CountPrefixed pfx f a) = CBLen pfx + CBLen (f a)

-- | The byte length of a count-prefixed type is the length of the prefix type
--   (holding the length of the type) plus the length of the type.
--
-- Bit confusing. How to explain this? TODO
instance (Prefix pfx, Foldable f, BLen pfx, BLen (f a))
  => BLen (CountPrefixed pfx f a) where
    blen :: CountPrefixed pfx f a -> Int
blen CountPrefixed pfx f a
rfa = pfx -> Int
forall a. BLen a => a -> Int
blen (forall a. Prefix a => Int -> a
lenToPfx @pfx (f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa)) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ f a -> Int
forall a. BLen a => a -> Int
blen f a
fa
      where fa :: f a
fa = CountPrefixed pfx f a -> f a
forall {k1} {k} (p :: k1) (f :: k -> Type) (x :: k).
Refined1 p f x -> f x
unrefine1 CountPrefixed pfx f a
rfa

instance (Prefix pfx, Foldable f, Put pfx, Put (f a))
  => Put (CountPrefixed pfx f a) where
    put :: CountPrefixed pfx f a -> Putter
put CountPrefixed pfx f a
rfa = pfx -> Putter
forall a. Put a => a -> Putter
put (forall a. Prefix a => Int -> a
lenToPfx @pfx (f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa)) Putter -> Putter -> Putter
forall a. Semigroup a => a -> a -> a
<> f a -> Putter
forall a. Put a => a -> Putter
put f a
fa
      where fa :: f a
fa = CountPrefixed pfx f a -> f a
forall {k1} {k} (p :: k1) (f :: k -> Type) (x :: k).
Refined1 p f x -> f x
unrefine1 CountPrefixed pfx f a
rfa

class GetCount f where getCount :: Get a => Int -> Getter (f a)
instance GetCount [] where getCount :: forall a. Get a => Int -> Getter [a]
getCount Int
n = Int -> ParserT PureMode E a -> ParserT PureMode E [a]
forall (m :: Type -> Type) a. Monad m => Int -> m a -> m [a]
Monad.count Int
n ParserT PureMode E a
forall a. Get a => Getter a
get

instance (Prefix pfx, GetCount f, Get pfx, Get a)
  => Get (CountPrefixed pfx f a) where
    get :: Getter (CountPrefixed pfx f a)
get = do
        pfx
pfx <- forall a. Get a => Getter a
get @pfx
        f a
fa <- Int -> Getter (f a)
forall a. Get a => Int -> Getter (f a)
forall (f :: Type -> Type) a.
(GetCount f, Get a) =>
Int -> Getter (f a)
getCount (pfx -> Int
forall a. Prefix a => a -> Int
pfxToLen pfx
pfx)
        CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a)
forall a. a -> ParserT PureMode E a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a))
-> CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a)
forall a b. (a -> b) -> a -> b
$ f a -> CountPrefixed pfx f a
forall {k} {k1} (f :: k -> Type) (x :: k) (p :: k1).
f x -> Refined1 p f x
reallyUnsafeRefine1 f a
fa