{-# LANGUAGE DataKinds, MultiParamTypeClasses, TypeFamilies #-}

-- | Module, containing some boilerplate for support of
-- arithmetic operations in Michelson language.

module Michelson.Typed.Arith
  ( ArithOp (..)
  , UnaryArithOp (..)
  , ArithError (..)
  , ArithErrorType (..)
  , Add
  , Sub
  , Mul
  , Abs
  , Neg
  , Or
  , And
  , Xor
  , Not
  , Lsl
  , Lsr
  , Compare
  , Eq'
  , Neq
  , Lt
  , Gt
  , Le
  , Ge
  , compareOp
  ) where

import Data.Bits (complement, shift, (.&.), (.|.))
import Data.Singletons (Sing, SingI(..))
import Fmt (Buildable(build))

import Michelson.Typed.Sing (SingT(..))
import Michelson.Typed.T (T(..))
import Michelson.Typed.Value (Comparable, Comparability(..), checkComparability, Value'(..))
import Tezos.Core (addMutez, mulMutez, subMutez, timestampFromSeconds, timestampToSeconds)

-- | Class for binary arithmetic operation.
--
-- Takes binary operation marker as @op@ parameter,
-- types of left operand @n@ and right operand @m@.
class ArithOp aop (n :: T) (m :: T) where

  -- | Type family @ArithRes@ denotes the type resulting from
  -- computing operation @op@ from operands of types @n@ and @m@.
  --
  -- For instance, adding integer to natural produces integer,
  -- which is reflected in following instance of type family:
  -- @ArithRes Add CNat CInt = CInt@.
  type ArithRes aop n m :: T

  -- | Evaluate arithmetic operation on given operands.
  evalOp
    :: proxy aop
    -> Value' instr n
    -> Value' instr m
    -> Either (ArithError (Value' instr n) (Value' instr m)) (Value' instr (ArithRes aop n m))

-- | Denotes the error type occured in the arithmetic operation.
data ArithErrorType
  = AddOverflow
  | MulOverflow
  | SubUnderflow
  | LslOverflow
  | LsrUnderflow
  deriving stock (Show, Eq, Ord, Generic)

instance NFData ArithErrorType

-- | Represents an arithmetic error of the operation.
data ArithError n m
  = MutezArithError ArithErrorType n m
  | ShiftArithError ArithErrorType n m
  deriving stock (Show, Eq, Ord, Generic)

instance (NFData n, NFData m) => NFData (ArithError n m)

-- | Marker data type for add operation.
class UnaryArithOp aop (n :: T) where
  type UnaryArithRes aop n :: T
  evalUnaryArithOp :: proxy aop -> Value' instr n -> Value' instr (UnaryArithRes aop n)

data Add
data Sub
data Mul
data Abs
data Neg

data Or
data And
data Xor
data Not
data Lsl
data Lsr

data Compare
data Eq'
data Neq
data Lt
data Gt
data Le
data Ge

instance ArithOp Add 'TNat 'TInt where
  type ArithRes Add 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i + j)
instance ArithOp Add 'TInt 'TNat where
  type ArithRes Add 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i + toInteger j)
instance ArithOp Add 'TNat 'TNat where
  type ArithRes Add 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i + j)
instance ArithOp Add 'TInt 'TInt where
  type ArithRes Add 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i + j)
instance ArithOp Add 'TTimestamp 'TInt where
  type ArithRes Add 'TTimestamp 'TInt = 'TTimestamp
  evalOp _ (VTimestamp i) (VInt j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds i + j
instance ArithOp Add 'TInt 'TTimestamp where
  type ArithRes Add 'TInt 'TTimestamp = 'TTimestamp
  evalOp _ (VInt i) (VTimestamp j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds j + i
instance ArithOp Add 'TMutez 'TMutez where
  type ArithRes Add 'TMutez 'TMutez = 'TMutez
  evalOp _ n@(VMutez i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError AddOverflow n m) (Right . VMutez) $ i `addMutez` j

instance ArithOp Sub 'TNat 'TInt where
  type ArithRes Sub 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i - j)
instance ArithOp Sub 'TInt 'TNat where
  type ArithRes Sub 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i - toInteger j)
instance ArithOp Sub 'TNat 'TNat where
  type ArithRes Sub 'TNat 'TNat = 'TInt
  evalOp _ (VNat i) (VNat j) = Right $ VInt (toInteger i - toInteger j)
instance ArithOp Sub 'TInt 'TInt where
  type ArithRes Sub 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i - j)
instance ArithOp Sub 'TTimestamp 'TInt where
  type ArithRes Sub 'TTimestamp 'TInt = 'TTimestamp
  evalOp _ (VTimestamp i) (VInt j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds i - j
instance ArithOp Sub 'TTimestamp 'TTimestamp where
  type ArithRes Sub 'TTimestamp 'TTimestamp = 'TInt
  evalOp _ (VTimestamp i) (VTimestamp j) =
    Right $ VInt $ timestampToSeconds i - timestampToSeconds j
instance ArithOp Sub 'TMutez 'TMutez where
  type ArithRes Sub 'TMutez 'TMutez = 'TMutez
  evalOp _ n@(VMutez i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError SubUnderflow n m) (Right . VMutez) $ i `subMutez` j

instance ArithOp Mul 'TNat 'TInt where
  type ArithRes Mul 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i * j)
instance ArithOp Mul 'TInt 'TNat where
  type ArithRes Mul 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i * toInteger j)
instance ArithOp Mul 'TNat 'TNat where
  type ArithRes Mul 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i * j)
instance ArithOp Mul 'TInt 'TInt where
  type ArithRes Mul 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i * j)
instance ArithOp Mul 'TNat 'TMutez where
  type ArithRes Mul 'TNat 'TMutez = 'TMutez
  evalOp _ n@(VNat i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError MulOverflow n m) (Right . VMutez) $ j `mulMutez` i
instance ArithOp Mul 'TMutez 'TNat where
  type ArithRes Mul 'TMutez 'TNat = 'TMutez
  evalOp _ n@(VMutez i) m@(VNat j) = res
    where
      res = maybe (Left $ MutezArithError MulOverflow n m) (Right . VMutez) $ i `mulMutez` j

instance UnaryArithOp Abs 'TInt where
  type UnaryArithRes Abs 'TInt = 'TNat
  evalUnaryArithOp _ (VInt i) = VNat (fromInteger $ abs i)

instance UnaryArithOp Neg 'TInt where
  type UnaryArithRes Neg 'TInt = 'TInt
  evalUnaryArithOp _ (VInt i) = VInt (-i)
instance UnaryArithOp Neg 'TNat where
  type UnaryArithRes Neg 'TNat = 'TInt
  evalUnaryArithOp _ (VNat i) = VInt (- fromIntegral i)

instance ArithOp Or 'TNat 'TNat where
  type ArithRes Or 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .|. j)
instance ArithOp Or 'TBool 'TBool where
  type ArithRes Or 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .|. j)

instance ArithOp And 'TInt 'TNat where
  type ArithRes And 'TInt 'TNat = 'TNat
  evalOp _ (VInt i) (VNat j) = Right $ VNat ((fromInteger i) .&. j)
instance ArithOp And 'TNat 'TNat where
  type ArithRes And 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .&. j)
instance ArithOp And 'TBool 'TBool where
  type ArithRes And 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .&. j)

instance ArithOp Xor 'TNat 'TNat where
  type ArithRes Xor 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i `xor` j)
instance ArithOp Xor 'TBool 'TBool where
  type ArithRes Xor 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i `xor` j)

instance ArithOp Lsl 'TNat 'TNat where
  type ArithRes Lsl 'TNat 'TNat = 'TNat
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LslOverflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (fromIntegral j))

instance ArithOp Lsr 'TNat 'TNat where
  type ArithRes Lsr 'TNat 'TNat = 'TNat
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LsrUnderflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (-(fromIntegral j)))

instance UnaryArithOp Not 'TInt where
  type UnaryArithRes Not 'TInt = 'TInt
  evalUnaryArithOp _ (VInt i) = VInt (complement i)
instance UnaryArithOp Not 'TNat where
  type UnaryArithRes Not 'TNat = 'TInt
  evalUnaryArithOp _ (VNat i) = VInt (complement $ toInteger i)
instance UnaryArithOp Not 'TBool where
  type UnaryArithRes Not 'TBool = 'TBool
  evalUnaryArithOp _ (VBool i) = VBool (not i)

compareOp :: forall t i. (Comparable t, SingI t) => Value' i t -> Value' i t -> Integer
compareOp a' b' = case (sing :: Sing t, a', b') of
  (STInt, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STNat, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STString, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STBytes, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STMutez, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STBool, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STKeyHash, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STTimestamp, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STAddress, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STPair l m, VPair (a, b), VPair (c, d)) ->
    case checkComparability l of
      CanBeCompared ->
        case compareOp a c of
          0  -> case checkComparability m of
            CanBeCompared -> compareOp b d
          r' -> r'

instance UnaryArithOp Eq' 'TInt where
  type UnaryArithRes Eq' 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i == 0)

instance UnaryArithOp Neq 'TInt where
  type UnaryArithRes Neq 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i /= 0)


instance UnaryArithOp Lt 'TInt where
  type UnaryArithRes Lt 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i < 0)

instance UnaryArithOp Gt 'TInt where
  type UnaryArithRes Gt 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i > 0)

instance UnaryArithOp Le 'TInt where
  type UnaryArithRes Le 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i <= 0)

instance UnaryArithOp Ge 'TInt where
  type UnaryArithRes Ge 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i >= 0)


instance Buildable ArithErrorType where
  build = \case
    AddOverflow -> "add overflow"
    MulOverflow -> "mul overflow"
    SubUnderflow -> "sub overflow"
    LslOverflow -> "lsl overflow"
    LsrUnderflow -> "lsr underflow"

instance (Show n, Show m) => Buildable (ArithError n m) where
  build (MutezArithError errType n m) = "Mutez "
    <> build errType <> " with " <> show n <> ", " <> show m
  build (ShiftArithError errType n m) =
    build errType <> " with " <> show n <> ", " <> show m