{-# LANGUAGE TypeOperators #-}

module Ribosome.Msgpack.Decode(
  MsgpackDecode(..),
) where

import Data.ByteString.Internal (unpackChars)
import Data.Map.Strict (Map, (!?))
import qualified Data.Map.Strict as Map (fromList, toList)
import Data.MessagePack (Object(..))
import GHC.Generics (
  Generic,
  Rep,
  (:*:)(..),
  M1(..),
  D1,
  C1,
  S1,
  K1(..),
  Selector,
  Constructor,
  selName,
  to,
  conIsRecord,
  )
import Ribosome.Msgpack.Util (Err)
import qualified Ribosome.Msgpack.Util as Util (string, invalid, missingRecordKey, illegalType)

class MsgpackDecode a where
  fromMsgpack :: Object -> Either Err a
  default fromMsgpack :: (Generic a, GMsgpackDecode (Rep a)) => Object -> Either Err a
  fromMsgpack = fmap to . gMsgpackDecode

  missingKey :: String -> Object -> Either Err a
  missingKey = Util.missingRecordKey

class GMsgpackDecode f where
  gMsgpackDecode :: Object -> Either Err (f a)

  gMissingKey :: String -> Object -> Either Err (f a)
  gMissingKey = Util.missingRecordKey

class MsgpackDecodeProd f where
  msgpackDecodeRecord :: Map Object Object -> Either Err (f a)
  msgpackDecodeProd :: [Object] -> Either Err ([Object], f a)

instance (GMsgpackDecode f) => GMsgpackDecode (D1 c f) where
  gMsgpackDecode =
    fmap M1 . gMsgpackDecode @f

instance (Constructor c, MsgpackDecodeProd f) => GMsgpackDecode (C1 c f) where
  gMsgpackDecode =
    fmap M1 . decode
    where
      isRec = conIsRecord (undefined :: t c f p)
      decode o@(ObjectMap om) =
        if isRec then msgpackDecodeRecord om else Util.invalid "illegal ObjectMap for product" o
      decode o | isRec =
        Util.invalid "illegal non-ObjectMap for record" o
      decode o =
        msgpackDecodeProd (prod o) >>= check
        where
          check ([], a) = Right a
          check _ = Util.invalid "too many values for product" o
          prod (ObjectArray oa) = oa
          prod ob = [ob]

instance (MsgpackDecodeProd f, MsgpackDecodeProd g) => MsgpackDecodeProd (f :*: g) where
  msgpackDecodeRecord o = do
    left <- msgpackDecodeRecord o
    right <- msgpackDecodeRecord o
    return $ left :*: right
  msgpackDecodeProd o = do
    (rest, left) <- msgpackDecodeProd o
    (rest1, right) <- msgpackDecodeProd rest
    return (rest1, left :*: right)

instance (Selector s, GMsgpackDecode f) => MsgpackDecodeProd (S1 s f) where
  msgpackDecodeRecord o =
    M1 <$> maybe (gMissingKey key (ObjectMap o)) gMsgpackDecode (o !? Util.string key)
    where
      key = selName (undefined :: t s f p)
  msgpackDecodeProd (cur:rest) = do
    a <- gMsgpackDecode cur
    return (rest, M1 a)
  msgpackDecodeProd [] = Util.invalid "too few values for product" ObjectNil

instance MsgpackDecode a => GMsgpackDecode (K1 i a) where
  gMsgpackDecode = fmap K1 . fromMsgpack

  gMissingKey key =
    fmap K1 . missingKey key

instance (Ord k, MsgpackDecode k, MsgpackDecode v) => MsgpackDecode (Map k v) where
  fromMsgpack (ObjectMap om) = do
    m <- traverse decodePair $ Map.toList om
    Right $ Map.fromList m
    where
      decodePair (k, v) = do
        k1 <- fromMsgpack k
        v1 <- fromMsgpack v
        return (k1, v1)
  fromMsgpack o = Util.illegalType "Map" o

instance MsgpackDecode Int where
  fromMsgpack (ObjectInt i) = Right $ fromIntegral i
  fromMsgpack o = Util.illegalType "Int" o

instance {-# OVERLAPPING #-} MsgpackDecode String where
  fromMsgpack (ObjectString os) = Right $ unpackChars os
  fromMsgpack o = Util.illegalType "String" o

instance {-# OVERLAPPABLE #-} MsgpackDecode a => MsgpackDecode [a] where
  fromMsgpack (ObjectArray oa) = traverse fromMsgpack oa
  fromMsgpack o = Util.illegalType "List" o

instance MsgpackDecode a => MsgpackDecode (Maybe a) where
  fromMsgpack ObjectNil = Right Nothing
  fromMsgpack o = Just <$> fromMsgpack o

  missingKey _ _ = Right Nothing

instance MsgpackDecode Bool where
  fromMsgpack (ObjectBool a) = Right a
  fromMsgpack o = Util.illegalType "Bool" o