module Michelson.TypeCheck.Value
    ( typeCheckValImpl
    , typeCheckCValue
    ) where

import Control.Monad.Except (liftEither, throwError)
import Data.Constraint (Dict(..))
import Data.Default (def)
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Singletons (SingI(..))
import Data.Typeable ((:~:)(..))
import Prelude hiding (EQ, GT, LT)

import Michelson.Text
import Michelson.TypeCheck.Error (TCError(..), TCTypeError(..))
import Michelson.TypeCheck.Helpers
import Michelson.TypeCheck.TypeCheck (TcInstrHandler, TypeCheckEnv(..), TypeCheckInstr)
import Michelson.TypeCheck.Types
import Michelson.Typed
  (pattern AsUType, CT(..), CValue(..), EpAddress(..), Notes(..), ParamNotes(..), Sing(..),
  Value'(..), converge, fromSingCT, fromSingT, starNotes)
import qualified Michelson.Typed as T
import qualified Michelson.Untyped as U
import Tezos.Address (Address(..))
import Tezos.Core
import Tezos.Crypto (parseKeyHash, parsePublicKey, parseSignature)

typeCheckCValue
  :: U.Value' op -> CT -> Either (U.Value' op, TCTypeError) SomeCValue
typeCheckCValue val ct = case (val, ct) of
  (U.ValueInt i, CInt) -> pure $ CvInt i :--: SCInt
  (U.ValueInt i, CNat)
    | i >= 0 -> pure $ CvNat (fromInteger i) :--: SCNat
    | otherwise -> Left (U.ValueInt i, NegativeNat)
  (v@(U.ValueInt i), CMutez) -> do
    mtz <- maybeToRight (v, InvalidTimestamp) . mkMutez $ fromInteger i
    pure $ CvMutez mtz :--: SCMutez
  (U.ValueString s, CString) ->
    pure $ CvString s :--: SCString
  (v@(U.ValueString s), CAddress) -> do
    addr <- T.parseEpAddress (unMText s) `onLeft` ((v, ) . InvalidAddress)
    pure $ CvAddress addr :--: SCAddress
  (v@(U.ValueString s), CKeyHash) -> do
    kHash <- parseKeyHash (unMText s) `onLeft` ((v, ) . InvalidKeyHash)
    pure $ CvKeyHash kHash :--: SCKeyHash
  (v@(U.ValueString s), CTimestamp) -> do
    tstamp <- maybeToRight (v, InvalidTimestamp) . parseTimestamp $ unMText s
    pure $ CvTimestamp tstamp :--: SCTimestamp
  (U.ValueInt i, CTimestamp) ->
    pure $ CvTimestamp (timestampFromSeconds i) :--: SCTimestamp
  (U.ValueBytes (U.InternalByteString s), CBytes) ->
    pure $ CvBytes s :--: SCBytes
  (U.ValueTrue, CBool) -> pure $ CvBool True :--: SCBool
  (U.ValueFalse, CBool) -> pure $ CvBool False :--: SCBool
  (v, t) ->
    Left $ (v, InvalidValueType (T.Tc t))

typeCheckCVals
  :: forall t op . (Typeable t, SingI t)
  => [U.Value' op]
  -> CT
  -> Either (U.Value' op, TCTypeError) [CValue t]
typeCheckCVals mvs t = traverse check mvs
  where
    check mv = do
      v :--: (_ :: Sing t') <- typeCheckCValue mv t
      Refl <- eqType @('T.Tc t) @('T.Tc t') `onLeft` (,) mv
      pure v

tcFailedOnValue :: U.Value -> T.T -> Text -> Maybe TCTypeError -> TypeCheckInstr a
tcFailedOnValue v t msg err = do
  loc <- ask
  throwError $ TCFailedOnValue v t msg loc err

-- | Function @typeCheckValImpl@ converts a single Michelson value
-- given in representation from @Michelson.Type@ module to representation
-- in strictly typed GADT.
--
-- As a third argument, @typeCheckValImpl@ accepts expected type of value.
--
-- Type checking algorithm pattern-matches on parse value representation,
-- expected type @t@ and constructs @Val t@ value.
--
-- If there was no match on a given pair of value and expected type,
-- that is interpreted as input of wrong type and type check finishes with
-- error.
typeCheckValImpl
  :: TcInstrHandler
  -> U.Value
  -> (Sing t, Notes t)
  -> TypeCheckInstr SomeNotedValue
typeCheckValImpl tcDo uvalue ty@(tySing, tyNotes) = case (uvalue, tySing, tyNotes) of
  (mv, t@(STc ct), NTc nt) -> do
    case typeCheckCValue mv (fromSingCT ct) of
      Left (uval, err) -> tcFailedOnValue uval (fromSingT $ t) "" (Just err)
      Right (v :--: cst) -> pure $ VC v :::: (STc cst, NTc nt)
  (U.ValueString (parsePublicKey . unMText -> Right s), STKey, _) ->
    pure $ T.VKey s :::: ty

  (U.ValueString (parseSignature . unMText -> Right s), STSignature, _) ->
    pure $ VSignature s :::: ty

  (U.ValueString (parseChainId . unMText -> Right ci), STChainId, _) ->
    pure $ VChainId ci :::: ty
  (U.ValueBytes (mkChainId . U.unInternalByteString -> Just ci), STChainId, _) ->
    pure $ VChainId ci :::: ty

  ( cv@(U.ValueString (T.parseEpAddress . unMText -> Right epAddr)),
                      STContract (pc :: Sing cp), NTContract _ pn ) -> do
    instrPos <- ask
    contracts <- gets tcContracts
    let ensureTypeMatches :: (Typeable t, SingI t) => (Sing t, Notes t) -> TypeCheckInstr (cp :~: t)
        ensureTypeMatches (_, pn') =
          liftEither @_ @TypeCheckInstr $
          first (TCFailedOnValue cv (fromSingT tySing) "wrong contract parameter" instrPos . Just) $
            fmap fst $ matchTypes pn pn'
    let unsupportedType :: Text -> Either TCError a
        unsupportedType desc =
          Left $
          TCFailedOnValue cv (fromSingT pc) (desc <> " in type argument of 'contract' type") instrPos $
          Just (UnsupportedTypes [fromSingT pc])
    let EpAddress addr epName = epAddr
    case addr of
      KeyAddress _ -> do
        Refl <- ensureTypeMatches (second unParamNotes T.tyImplicitAccountParam)
        pure $ VContract addr T.sepcPrimitive :::: ty
      ContractAddress ca ->
        case M.lookup ca contracts of
          -- Wrapping into 'ParamNotesUnsafe' is safe because originated contract has
          -- valid parameter type
          Just (AsUType cpSing (T.ParamNotesUnsafe -> cpNotes)) -> do
            Dict <- liftEither $ maybe (unsupportedType "Operation") pure (T.opAbsense cpSing)
            Dict <- liftEither $ maybe (unsupportedType "Nested BigMaps") pure (T.nestedBigMapsAbsense cpSing)
            case T.mkEntryPointCall epName (cpSing, cpNotes) of
              Nothing ->
                throwError $
                TCFailedOnValue cv (fromSingT tySing) "unknown entrypoint" instrPos . Just $
                EntryPointNotFound epName
              Just (T.MkEntryPointCallRes argNotes epc) -> do
                ensureTypeMatches (sing, argNotes)
                pure $ VContract addr (T.SomeEpc epc) :::: (sing, NTContract U.noAnn argNotes)
          Nothing ->
            throwError $ TCFailedOnValue cv (fromSingT tySing) "Contract literal unknown"
              instrPos (Just $ UnknownContract addr)

  (U.ValueUnit, STUnit, _) -> pure $ VUnit :::: ty
  (U.ValuePair ml mr, STPair lt rt, NTPair n1 n2 n3 nl nr) -> do
    l :::: (lst, ln) <- typeCheckValImpl tcDo ml (lt, nl)
    r :::: (rst, rn) <- typeCheckValImpl tcDo mr (rt, nr)
    pure $ VPair (l, r) :::: (STPair lst rst, NTPair n1 n2 n3 ln rn)
  (U.ValueLeft ml, STOr lt rt, NTOr n1 n2 n3 nl nr) -> do
    l :::: (lst, ln) <- typeCheckValImpl tcDo ml (lt, nl)
    pure $ VOr (Left l) :::: ( STOr lst rt, NTOr n1 n2 n3 ln nr)
  (U.ValueRight mr, STOr lt rt, NTOr n1 n2 n3 nl nr) -> do
    r :::: (rst, rn) <- typeCheckValImpl tcDo mr (rt, nr)
    pure $ VOr (Right r) :::: ( STOr lt rst, NTOr n1 n2 n3 nl rn)
  (U.ValueSome mv, STOption vt, NTOption na nt) -> do
    v :::: (vst, vns) <- typeCheckValImpl tcDo mv (vt, nt)
    pure $ VOption (Just v) :::: (STOption vst, NTOption na vns)
  (U.ValueNone, STOption _, _) -> do
    pure $ VOption Nothing :::: ty

  (U.ValueNil, STList _, _) ->
    pure $ T.VList [] :::: ty

  (U.ValueSeq (toList -> mels), STList vt, NTList _ x) -> do
    (els, _) <- typeCheckValsImpl tcDo mels (vt, x)
    pure $ VList els :::: ty

  (U.ValueNil, STSet _, _) ->
    pure $ T.VSet S.empty :::: ty

  (sq@(U.ValueSeq (toList -> mels)), STSet vt, _) -> do
    instrPos <- ask
    els <- liftEither $ typeCheckCVals mels (fromSingCT vt)
            `onLeft` \(cv, err) -> TCFailedOnValue cv (fromSingT $ STc vt)
                                        "wrong type of set element:" instrPos (Just err)
    elsS <- liftEither $ S.fromDistinctAscList <$> ensureDistinctAsc id els
              `onLeft` \msg -> TCFailedOnValue sq (fromSingT $ STc vt) msg instrPos Nothing
    pure $ VSet elsS :::: ty

  (U.ValueNil, STMap _ _, _) -> pure $ T.VMap M.empty :::: ty

  (sq@(U.ValueMap (toList -> mels)), STMap kt vt, NTMap _ _ vn) -> do
    keyOrderedElts <- typeCheckMapVal tcDo mels sq vn kt vt
    pure $ VMap (M.fromDistinctAscList keyOrderedElts) :::: ty

  (U.ValueNil, STBigMap _ _ , _) ->
    pure $ T.VBigMap M.empty :::: ty

  (sq@(U.ValueMap (toList -> mels)), STBigMap kt vt, NTBigMap _ _ vn) -> do
    keyOrderedElts <- typeCheckMapVal tcDo mels sq vn kt vt
    pure $ VBigMap (M.fromDistinctAscList keyOrderedElts) :::: ty

  (v, STLambda (it :: Sing it) (ot :: Sing ot), NTLambda vn _ _) -> do
    mp <- case v of
      U.ValueNil       -> pure []
      U.ValueLambda mp -> pure $ toList mp
      _ -> tcFailedOnValue v (fromSingT tySing) "unexpected value" Nothing
    li :/ instr <- typeCheckImpl tcDo mp ((it, starNotes, def) ::& SNil)
    let (_, ins, _) ::& SNil = li
    let lamS = STLambda it ot
    case instr of
      lam ::: (lo :: HST lo) -> do
        case eqHST1 @ot lo of
          Right Refl -> do
            let (_, ons, _) ::& SNil = lo
            pure $ VLam (T.RfNormal lam) :::: (STLambda it ot, NTLambda vn ins ons)
          Left m ->
            tcFailedOnValue v (fromSingT tySing)
                    "wrong output type of lambda's value:" (Just m)
      AnyOutInstr lam ->
        pure $ VLam (T.RfAlwaysFails lam) :::: (lamS, NTLambda def ins starNotes)

  (v, t, _) -> tcFailedOnValue v (fromSingT t) "unknown value" Nothing

-- | Function @typeCheckMapVal@ typechecks given list of @Elt@s and
-- ensures, that its keys are in ascending order.
--
-- It return list of pairs (key, value) with keys in ascending order
-- so it is safe to call @fromDistinctAscList@ on returned list
typeCheckMapVal
  :: (SingI kt, Typeable kt, SingI vt, Typeable vt)
  => TcInstrHandler
  -> [U.Elt U.ExpandedOp]
  -> U.Value
  -> Notes vt
  -> Sing kt
  -> Sing vt
  -> TypeCheckInstr [(CValue kt, T.Value vt)]
typeCheckMapVal tcDo mels sq vn kt vt = do
  instrPos <- ask
  ks <- liftEither $ typeCheckCVals (map (\(U.Elt k _) -> k) mels) (fromSingCT kt)
          `onLeft` \(cv, err) -> TCFailedOnValue cv (fromSingT $ STc kt)
                                      "wrong type of map key:" instrPos (Just err)
  (vals, _) <- typeCheckValsImpl tcDo (map (\(U.Elt _ v) -> v) mels) (vt, vn)
  ksS <- liftEither $ ensureDistinctAsc id ks
        `onLeft` \msg -> TCFailedOnValue sq (fromSingT $ STc kt) msg instrPos Nothing
  pure $ zip ksS vals

typeCheckValsImpl
  :: forall t . (Typeable t, SingI t)
  => TcInstrHandler
  -> [U.Value]
  -> (Sing t, Notes t)
  -> TypeCheckInstr ([T.Value t], Notes t)
typeCheckValsImpl tcDo mvs (t, nt) =
  fmap (first reverse) $ foldM check ([], nt) mvs
  where
    check (res, ns) mv = do
      instrPos <- ask
      v :::: ((_ :: Sing t'), vns) <- typeCheckValImpl tcDo mv (t, nt)
      Refl <- liftEither $ eqType @t @t' `onLeft`
        (TCFailedOnValue mv (fromSingT t) "wrong element type" instrPos . Just)
      ns' <- liftEither $ converge ns vns `onLeft`
        ((TCFailedOnValue mv (fromSingT t) "wrong element type") instrPos . Just . AnnError)
      pure (v : res, ns')