-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Module, containing type classes for operating with Michelson values
-- in the context of polymorphic stack type operations.

module Michelson.Typed.Polymorphic
  ( EDivOp (..)
  , MemOp (..)
  , MapOp (..)
  , IterOp (..)
  , SizeOp (..)
  , GetOp (..)
  , UpdOp (..)
  , SliceOp (..)
  , ConcatOp (..)
  , divMich
  , modMich
  ) where

import qualified Data.ByteString as B
import qualified Data.Map as M
import qualified Data.Set as S

import Michelson.Text
import Michelson.Typed.Annotation
import Michelson.Typed.Sing (KnownT)
import Michelson.Typed.T (T(..))
import Michelson.Typed.Value (Value'(..))
import Michelson.Untyped.Annotation (noAnn)

import Tezos.Core (divModMutez, divModMutezInt)

class MemOp (c :: T) where
  type MemOpKey c :: T
  evalMem :: Value' instr (MemOpKey c) -> Value' instr c -> Bool
instance MemOp ('TSet e) where
  type MemOpKey ('TSet e) = e
  evalMem e (VSet s) = e `S.member` s
instance MemOp ('TMap k v) where
  type MemOpKey ('TMap k v) = k
  evalMem k (VMap m) = k `M.member` m
instance MemOp ('TBigMap k v) where
  type MemOpKey ('TBigMap k v) = k
  evalMem k (VBigMap m) = k `M.member` m

class MapOp (c :: T) where
  type MapOpInp c :: T
  type MapOpRes c :: T -> T
  mapOpToList :: Value' instr c -> [Value' instr (MapOpInp c)]
  mapOpFromList
    :: (KnownT b)
    => Value' instr c -> [Value' instr b] -> Value' instr (MapOpRes c b)
instance MapOp ('TMap k v) where
  type MapOpInp ('TMap k v) = 'TPair k v
  type MapOpRes ('TMap k v) = 'TMap k
  mapOpToList (VMap m) = map (\(k, v) -> VPair (k, v)) $ M.toAscList m
  mapOpFromList (VMap m) l =
    VMap $ M.fromList $ zip (map fst $ M.toAscList m) l
instance MapOp ('TList e) where
  type MapOpInp ('TList e) = e
  type MapOpRes ('TList e) = 'TList
  mapOpToList (VList l) = l
  mapOpFromList (VList _) l' = VList l'
-- If you find it difficult to implement 'MapOp' for your datatype
-- because of order of type arguments in it, consider wrapping it
-- into a newtype.

class IterOp (c :: T) where
  type IterOpEl c :: T
  iterOpDetachOne ::
    Value' instr c -> (Maybe (Value' instr (IterOpEl c)), Value' instr c)
instance IterOp ('TMap k v) where
  type IterOpEl ('TMap k v) = 'TPair k v
  iterOpDetachOne (VMap m) =
    (VPair <$> M.lookupMin m, VMap $ M.deleteMin m)
instance IterOp ('TList e) where
  type IterOpEl ('TList e) = e
  iterOpDetachOne (VList l) =
    case l of
      x : xs -> (Just x, VList xs)
      [] -> (Nothing, VList [])
instance IterOp ('TSet e) where
  type IterOpEl ('TSet e) = e
  iterOpDetachOne (VSet s) = (S.lookupMin s, VSet $ S.deleteMin s)

class SizeOp (c :: T) where
  evalSize :: Value' instr c -> Int
instance SizeOp 'TString where
  evalSize (VString s) = length s
instance SizeOp ('TBytes) where
  evalSize (VBytes b) = length b
instance SizeOp ('TSet a) where
  evalSize (VSet s) = S.size s
instance SizeOp ('TList a) where
  evalSize (VList l) = length l
instance SizeOp ('TMap k v) where
  evalSize (VMap m) = M.size m

class UpdOp (c :: T) where
  type UpdOpKey c :: T
  type UpdOpParams c :: T
  evalUpd
    :: Value' instr (UpdOpKey c)
    -> Value' instr (UpdOpParams c) -> Value' instr c -> Value' instr c
instance UpdOp ('TMap k v) where
  type UpdOpKey ('TMap k v) = k
  type UpdOpParams ('TMap k v) = 'TOption v
  evalUpd k (VOption o) (VMap m) =
    case o of
      Just newV -> VMap $ M.insert k newV m
      Nothing -> VMap $ M.delete k m
instance UpdOp ('TBigMap k v) where
  type UpdOpKey ('TBigMap k v) = k
  type UpdOpParams ('TBigMap k v) = 'TOption v
  evalUpd k (VOption o) (VBigMap m) =
    case o of
      Just newV -> VBigMap $ M.insert k newV m
      Nothing -> VBigMap $ M.delete k m
instance UpdOp ('TSet a) where
  type UpdOpKey ('TSet a) = a
  type UpdOpParams ('TSet a) = 'TBool
  evalUpd k (VBool b) (VSet s) =
    case b of
      True -> VSet $ S.insert k s
      False -> VSet $ S.delete k s

class GetOp (c :: T) where
  type GetOpKey c :: T
  type GetOpVal c :: T
  evalGet :: Value' instr (GetOpKey c) -> Value' instr c -> Maybe (Value' instr (GetOpVal c))
instance GetOp ('TBigMap k v) where
  type GetOpKey ('TBigMap k v) = k
  type GetOpVal ('TBigMap k v) = v
  evalGet k (VBigMap m) = k `M.lookup` m
instance GetOp ('TMap k v) where
  type GetOpKey ('TMap k v) = k
  type GetOpVal ('TMap k v) = v
  evalGet k (VMap m) = k `M.lookup` m

class ConcatOp (c :: T) where
  evalConcat :: Value' instr c -> Value' instr c -> Value' instr c
  evalConcat' :: [Value' instr c] -> Value' instr c
instance ConcatOp ('TString) where
  evalConcat (VString s1) (VString s2) = (VString) (s1 <> s2)
  evalConcat' l =
    (VString) $ mconcat $ map (\(VString s) -> s) l
instance ConcatOp ('TBytes) where
  evalConcat (VBytes b1) (VBytes b2) = VBytes (b1 <> b2)
  evalConcat' l =
    (VBytes) $ foldr ((<>) . (\(VBytes b) -> b)) mempty l

class SliceOp (c :: T) where
  evalSlice :: Natural -> Natural -> Value' instr c -> Maybe (Value' instr c)
instance SliceOp 'TString where
  evalSlice o l (VString s) =
    VString <$> sliceImpl dropMText takeMText o l s
instance SliceOp 'TBytes where
  evalSlice o l (VBytes b) =
    VBytes <$> sliceImpl B.drop B.take o l b

sliceImpl ::
  Container str
  => (Int -> str -> str)
  -> (Int -> str -> str)
  -> Natural
  -> Natural
  -> str
  -> Maybe str
sliceImpl dropF takeF offset l s
  | offset >= fromIntegral (length s) || offset + l > fromIntegral (length s) =
    Nothing
  | otherwise
  -- Drop offset and then take requested number of items.
   = Just . takeF (fromIntegral l) . dropF (fromIntegral offset) $ s

class EDivOp (n :: T) (m :: T) where
  type EDivOpRes n m :: T
  type EModOpRes n m :: T

  -- | Converge the notes of given operands.
  convergeEDiv
    :: Notes n
    -> Notes m
    -> Either AnnConvergeError
        (Notes ('TOption ('TPair (EDivOpRes n m)
                                  (EModOpRes n m))))

  evalEDivOp
    :: Value' instr n
    -> Value' instr m
    -> Value' instr ('TOption ('TPair (EDivOpRes n m)
                                     (EModOpRes n m)))

instance EDivOp 'TInt 'TInt where
  type EDivOpRes 'TInt 'TInt = 'TInt
  type EModOpRes 'TInt 'TInt = 'TNat
  convergeEDiv n1 n2 =
    (\a -> NTOption noAnn $ NTPair noAnn noAnn noAnn a $ NTNat noAnn)
      <$> converge n1 n2
  evalEDivOp (VInt i) (VInt j) =
    if j == 0
      then VOption $ Nothing
      else VOption $ Just $
        VPair (VInt (divMich i j), VNat $ fromInteger $ modMich i j)
instance EDivOp 'TInt 'TNat where
  type EDivOpRes 'TInt 'TNat = 'TInt
  type EModOpRes 'TInt 'TNat = 'TNat
  convergeEDiv n1 _ = Right $ NTOption noAnn $ NTPair noAnn noAnn noAnn n1
    $ NTNat noAnn
  evalEDivOp (VInt i) (VNat j) =
    if j == 0
      then VOption $ Nothing
      else VOption $ Just $
        VPair (VInt (divMich i (toInteger j)), VNat $ fromInteger $ modMich i (toInteger j))
instance EDivOp 'TNat 'TInt where
  type EDivOpRes 'TNat 'TInt = 'TInt
  type EModOpRes 'TNat 'TInt = 'TNat
  convergeEDiv n1 _ = Right $ NTOption noAnn $ NTPair noAnn noAnn noAnn
    (NTInt noAnn) n1
  evalEDivOp (VNat i) (VInt j) =
    if j == 0
      then VOption $ Nothing
      else VOption $ Just $
        VPair (VInt (divMich (toInteger i) j), VNat $ fromInteger $ modMich (toInteger i) j)
instance EDivOp 'TNat 'TNat where
  type EDivOpRes 'TNat 'TNat = 'TNat
  type EModOpRes 'TNat 'TNat = 'TNat
  convergeEDiv n1 n2 = (\a -> NTOption noAnn $ NTPair noAnn noAnn noAnn a a)
    <$> converge n1 n2
  evalEDivOp (VNat i) (VNat j) =
    if j == 0
      then VOption $ Nothing
      else VOption $ Just $
        VPair (VNat (divMich i j), VNat $ (modMich i j))
instance EDivOp 'TMutez 'TMutez where
  type EDivOpRes 'TMutez 'TMutez = 'TNat
  type EModOpRes 'TMutez 'TMutez = 'TMutez
  convergeEDiv n1 n2 = (\a -> NTOption noAnn $ NTPair noAnn noAnn noAnn (NTNat noAnn) a)
    <$> converge n1 n2
  evalEDivOp (VMutez i) (VMutez j) =
    VOption $
    i `divModMutez` j <&> \case
      (quotient, remainder) ->
        VPair (VNat (fromIntegral quotient), VMutez remainder)

instance EDivOp 'TMutez 'TNat where
  type EDivOpRes 'TMutez 'TNat = 'TMutez
  type EModOpRes 'TMutez 'TNat = 'TMutez
  convergeEDiv n1 _ = Right $ NTOption noAnn $ NTPair noAnn noAnn noAnn n1 n1
  evalEDivOp (VMutez i) (VNat j) =
    VOption $
    i `divModMutezInt` j <&> \case
      (quotient, remainder) ->
        VPair (VMutez quotient, VMutez remainder)

-- | Computing 'div' function in Michelson style.
-- When divisor is negative, Haskell gives x as integer part,
-- while Michelson gives x+1.
divMich :: Integral a => a -> a -> a
divMich divisible divisor = divisible `div` divisor + extra
  where
    extra =
      if divisor > 0 ||
         divisible `mod` divisor == 0
      then 0
      else 1

-- | Computing 'mod' function in Michelson style.
-- When divisor is negative, Haskell gives a negative modulo,
-- while there is a positive modulo in Michelson.
modMich :: Integral a => a -> a -> a
modMich divisible divisor = divisible - divisor * intPart
  where intPart = divMich divisible divisor