{-# Language  TemplateHaskell #-}
{-# Language  DeriveFunctor #-}
module HBS2.Merkle where

import HBS2.Prelude
import HBS2.Hash

import Control.Applicative
import Codec.Serialise (serialise, deserialiseOrFail)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.Data
import Data.Foldable (forM_, traverse_)
import Data.List qualified as List
import Data.Text (Text)
import GHC.Generics
import Lens.Micro.Platform
import Control.Monad.Trans.Maybe
import Control.Monad
import Prettyprinter

newtype MerkleHash = MerkleHash { fromMerkleHash :: Hash HbSync }
                     deriving newtype (Eq,Ord,IsString,Pretty)
                     deriving stock (Data,Generic)

class Foldable t => Chunks t a where
  chunks :: Int -> a -> t a

instance Chunks [] ByteString where
  chunks i bs = go [] bs
    where
      size = fromIntegral i
      go acc x | BS.null x  = acc
               | BS.length x <= size = acc <> [x]
               | otherwise = go (acc <> [BS.take size x]) (BS.drop size x)

instance Chunks [] LBS.ByteString where
  chunks i bs = go [] bs
    where
      size = fromIntegral i
      go acc x | LBS.null x  = acc
               | LBS.length x <= size = acc <> [x]
               | otherwise = go (acc <> [LBS.take size x]) (LBS.drop size x)

instance Chunks [] [a] where
  chunks i xs = go xs
    where
      go [] = []
      go es | length es <= i = [es]
            | otherwise = let (p, ps) = List.splitAt i es in p : go ps


data PTree  a = T [PTree a] | L a
                deriving stock (Show, Functor, Generic)


instance Foldable PTree where
  foldr fn acc (L a)  = fn a acc
  foldr fn acc (T xs) = go acc xs
    where
      go b [] = b
      go b (y:ys) = foldr fn (go b ys) y

instance Traversable PTree where
  traverse fn (L a) = L <$> fn a
  traverse fn (T xs) = T <$> traverse (traverse fn) xs

newtype MaxNum a = MaxNum a
newtype MaxSize a = MaxSize a


newtype MNodeData =
  MNodeData
  { _mnodeHeight :: Integer
  }
  deriving stock (Generic,Data,Show)

makeLenses ''MNodeData

instance Serialise MNodeData

data AnnMetaData = NoMetaData | ShortMetadata Text | AnnHashRef (Hash HbSync)
  deriving stock (Generic,Data,Show,Eq)

instance Serialise AnnMetaData

data MTreeAnn a = MTreeAnn
 { _mtaMeta :: !AnnMetaData
 , _mtaCrypt :: !MTreeEncryption
 , _mtaTree :: !(MTree a)
 }
 deriving stock (Generic,Data,Show)

instance Serialise a => Serialise (MTreeAnn a)

data MerkleEncryptionType
  deriving stock (Data)

data MTreeEncryption
  = NullEncryption
  | CryptAccessKeyNaClAsymm (Hash HbSync)
  | EncryptGroupNaClSymm (Hash HbSync) ByteString
  deriving stock (Eq,Generic,Data,Show)

instance Serialise MTreeEncryption

data MTree a = MNode MNodeData [Hash HbSync] | MLeaf a
               deriving stock (Generic,Data,Show)

instance Serialise a => Serialise (MTree a)

newMNode :: Integer -> [Hash HbSync] -> MTree a
newMNode h = MNode (MNodeData h)

toPTree :: (Chunks [] a)
        => MaxSize Int
        -> MaxNum Int
        -> a
        -> PTree a

toPTree (MaxSize s) (MaxNum n) items | n <= 1 =
  case T $ fmap L (chunks s items) of
    T [L x] -> L x
    _       -> T []

toPTree (MaxSize s) (MaxNum n) items = go $ T (fmap L (chunks s items))
  where
    go (T []) = T []
    go (T [L x]) = L x
    go (T xs) | length xs <= n = T xs
              | otherwise = go $ T $ fmap go [ T x | x <- chunks n xs ]

    go leaf = leaf


makeMerkle :: (Monad m, Serialise a, Serialise (MTree a))
           => Integer -- | initial height
           -> PTree a
           -> ((Hash HbSync, MTree a, LBS.ByteString) -> m ())
           -> m (Hash HbSync)

makeMerkle h0 pt f = fst <$> go h0 pt
  where
    go hx (T xs) = do
      rs <- mapM (go hx) xs
      let hxx = maximumDef hx (fmap snd rs)
      let o = newMNode hxx (fmap fst rs)
      let bs = serialise o
      let h  = hashObject bs
      f (h, o, bs)
      pure (h, 1+hxx)

    go hx (L x) = do
      let o = MLeaf x
      let bs = serialise o
      let h  = hashObject bs
      f (h, o, bs)
      pure (h, 1+hx)



-- NOTE: MerkleAnn-note
--   MerkleAnn ломает всю красоту, но, видимо,
--   с этим уже ничего не поделать, пусть живёт.
--   Надо было аннотации просто класть 0-ым блоком
--   в меркл дерево, или вообще делать cons-ячейки,
--   но что уж теперь.

walkMerkle' :: forall a m . (Serialise (MTree a), (Serialise (MTreeAnn a)), Monad m)
           => Hash HbSync
           -> ( Hash HbSync -> m (Maybe LBS.ByteString) )
           -> ( Either (Hash HbSync) (MTree a) -> m () )
           -> m ()

walkMerkle' root flookup sink = go root
  where
    go hash = void $ runMaybeT do

      bs0 <- lift $ flookup hash

      bs <- MaybeT $ maybe1 bs0 (sink (Left hash) >> pure Nothing) (pure . Just)

      let t1 = deserialiseOrFail @(MTree a) bs

      either (const $ runWithAnnTree hash bs) runWithTree t1

    runWithAnnTree hash bs  = do
      let t = deserialiseOrFail @(MTreeAnn a) bs
      case t of
        Left{} -> lift (sink (Left hash)) >> mzero
        Right (MTreeAnn { _mtaTree = t1 }) -> runWithTree t1

    runWithTree t = lift do
      case t of
        n@(MLeaf _) -> sink (Right n)
        n@(MNode _ hashes) -> sink (Right n) >> traverse_ go hashes

walkMerkle :: forall a m . (Serialise (MTree a), Serialise (MTreeAnn a), Serialise a, Monad m)
           => Hash HbSync
           -> ( Hash HbSync -> m (Maybe LBS.ByteString) )
           -> ( Either (Hash HbSync) a -> m () )
           -> m ()

walkMerkle root flookup sink = do

  walkMerkle' root flookup withTree

  where
    withTree = \case
        (Right (MLeaf s))   -> sink (Right s)
        (Right (MNode _ _)) -> pure ()
        Left hx             -> sink (Left hx)


walkMerkleTree :: (Serialise (MTree a), Serialise a, Monad m)
           => MTree a
           -> ( Hash HbSync -> m (Maybe LBS.ByteString) )
           -> ( Either (Hash HbSync) a -> m () )
           -> m ()

walkMerkleTree tree flookup sink = case tree of
    (MLeaf s)   -> sink (Right s)
    (MNode _ hashes) -> forM_ hashes \h -> walkMerkle h flookup sink