module Michelson.TypeCheck.Value
    ( typeCheckValImpl
    ) 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 (Sing, SingI(..), demote)
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 AsUTypeExt, EpAddress(..), Notes(..), ParamNotes(..), SingT(..), Value'(..), 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)

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.
--
-- @typeCheckValImpl@ is polymorphic in the 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
  :: forall ty.
     SingI ty
  => TcInstrHandler
  -> U.Value
  -> TypeCheckInstr (T.Value ty)
typeCheckValImpl tcDo uvalue = case (uvalue, sing @ty) of
  (U.ValueInt i, STInt) -> pure $ T.VInt i
  (v@(U.ValueInt i), t@STNat)
    | i >= 0 -> pure $ VNat (fromInteger i)
    | otherwise -> tcFailedOnValue v (fromSingT t) "" (Just NegativeNat)
  (v@(U.ValueInt i), t@STMutez) -> do
    case mkMutez $ fromInteger i of
      Just mtz -> pure $ VMutez mtz
      Nothing -> tcFailedOnValue v (fromSingT t) "" (Just InvalidTimestamp)
  (U.ValueString s, STString) ->
    pure $ VString s
  (v@(U.ValueString s), t@STAddress) -> do
    case T.parseEpAddress (unMText s) of
      Right addr -> pure $ VAddress addr
      Left err -> tcFailedOnValue v (fromSingT t) "" (Just $ InvalidAddress err)
  (v@(U.ValueString s), t@STKeyHash) -> do
    case parseKeyHash (unMText s)  of
      Right kHash -> pure $ VKeyHash kHash
      Left err -> tcFailedOnValue v (fromSingT t) "" (Just $ InvalidKeyHash err)
  (v@(U.ValueString s), t@STTimestamp) -> do
    case parseTimestamp $ unMText s of
      Just tstamp -> pure $ VTimestamp tstamp
      Nothing -> tcFailedOnValue v (fromSingT t) "" (Just InvalidTimestamp)
  (U.ValueInt i, STTimestamp) ->
    pure $ VTimestamp (timestampFromSeconds i)
  (U.ValueBytes (U.InternalByteString s), STBytes) ->
    pure $ VBytes s
  (U.ValueTrue, STBool) -> pure $ VBool True
  (U.ValueFalse, STBool) -> pure $ VBool False

  (U.ValueString (parsePublicKey . unMText -> Right s), STKey) ->
    pure $ T.VKey s

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

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

  ( cv@(U.ValueString (T.parseEpAddress . unMText -> Right epAddr)),
                      STContract (pc :: Sing cp)) -> do
    instrPos <- ask
    contracts <- gets tcContracts
    let ensureTypeMatches :: forall t'. (Typeable t', SingI t') => TypeCheckInstr (cp :~: t')
        ensureTypeMatches =
          liftEither @_ @TypeCheckInstr $
          first (TCFailedOnValue cv (demote @ty) "wrong contract parameter" instrPos . Just) $
            eqType @cp @t'
    let unsupportedType :: T.BadTypeForScope -> TCError
        unsupportedType reason =
          TCFailedOnValue cv (fromSingT pc) ("Unsupported type in type argument of 'contract' type") instrPos $
          Just $ UnsupportedTypeForScope (fromSingT pc) reason
    let EpAddress addr epName = epAddr
    case addr of
      KeyAddress _ -> do
        Refl <- ensureTypeMatches @'T.TUnit
        pure $ VContract addr T.sepcPrimitive
      ContractAddress ca ->
        case M.lookup ca contracts of
          Just (U.ParameterType (AsUTypeExt (_ :: Sing cp') cpNotes) rootAnn) -> do
            -- Wrapping into 'ParamNotesUnsafe' is safe because originated contract has
            -- valid parameter type
            let paramNotes = ParamNotesUnsafe { pnNotes = cpNotes, pnRootAnn = rootAnn }
            Dict <- liftEither $ T.checkScope @(T.ParameterScope cp')
              `onLeft` unsupportedType
            case T.mkEntryPointCall epName paramNotes of
              Nothing ->
                throwError $
                TCFailedOnValue cv (demote @ty) "unknown entrypoint" instrPos . Just $
                EntryPointNotFound epName
              Just (T.MkEntryPointCallRes (_ :: Notes t') epc) -> do
                Refl <- ensureTypeMatches @t'
                pure $ VContract addr (T.SomeEpc epc)
          Nothing ->
            throwError $ TCFailedOnValue cv (demote @ty) "Contract literal unknown"
              instrPos (Just $ UnknownContract addr)

  (U.ValueUnit, STUnit) -> pure $ VUnit
  (U.ValuePair ml mr, STPair _ _) -> do
    l <- typeCheckValImpl tcDo ml
    r <- typeCheckValImpl tcDo mr
    pure $ VPair (l, r)
  (U.ValueLeft ml, STOr{}) -> do
    l <- typeCheckValImpl tcDo ml
    pure $ VOr (Left l)
  (U.ValueRight mr, STOr{}) -> do
    r <- typeCheckValImpl tcDo mr
    pure $ VOr (Right r)
  (U.ValueSome mv, STOption{}) -> do
    v <- typeCheckValImpl tcDo mv
    pure $ VOption (Just v)
  (U.ValueNone, STOption _) -> do
    pure $ VOption Nothing

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

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

  (U.ValueNil, STSet (s :: Sing st)) -> do
    instrPos <- ask
    case T.getComparableProofS s of
      Just Dict -> pure (T.VSet S.empty)
      Nothing -> throwError $ TCFailedOnValue uvalue (demote @st) "Non comparable types are not allowed in Sets"
        instrPos (Just $ UnsupportedTypeForScope (demote @st) T.BtNotComparable)

  (sq@(U.ValueSeq (toList -> mels)), s@(STSet (vt :: Sing st))) -> withComparable vt sq s $ do
    instrPos <- ask

    els <- typeCheckValsImpl tcDo mels
    elsS <- liftEither $ S.fromDistinctAscList <$> ensureDistinctAsc id els
              `onLeft` \msg -> TCFailedOnValue sq (fromSingT vt) msg instrPos Nothing
    pure $ VSet elsS

  (v@U.ValueNil, s@(STMap (st :: Sing st) _)) -> withComparable st v s $ pure $ T.VMap M.empty

  (sq@(U.ValueMap (toList -> mels)), s@(STMap (kt :: Sing st) _)) -> withComparable kt sq s $ do
    keyOrderedElts <- typeCheckMapVal tcDo mels sq kt
    pure $ VMap (M.fromDistinctAscList keyOrderedElts)

  (v@U.ValueNil, s@(STBigMap (st :: Sing st) _)) ->
    withComparable st v s $ pure $ T.VBigMap M.empty

  (sq@(U.ValueMap (toList -> mels)), s@(STBigMap (kt :: Sing st) _)) -> withComparable kt sq s $ do
    keyOrderedElts <- typeCheckMapVal tcDo mels sq kt
    pure $ VBigMap (M.fromDistinctAscList keyOrderedElts)

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

  (v, t) -> tcFailedOnValue v (fromSingT t) "unknown value" Nothing
  where
    withWTP :: forall t a. SingI t => (T.WellTyped t => TypeCheckInstr a) -> TypeCheckInstr a
    withWTP fn = case getWTP @t of
      Just Dict -> fn
      Nothing -> tcFailedOnValue uvalue (fromSingT $ sing @ty) "Value is not well typed" Nothing

withComparable
  :: forall a (t :: T.T) ty. Sing a
  -> U.Value
  -> Sing t
  -> (T.Comparable a => TypeCheckInstr ty)
  -> TypeCheckInstr ty
withComparable s uv t act = case T.getComparableProofS s of
  Just Dict -> act
  Nothing -> do
    instrPos <- ask
    liftEither $ Left $ TCFailedOnValue uv (fromSingT t) "Require a comparable type here" instrPos 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
  :: forall kt vt. (SingI kt, SingI vt)
  => TcInstrHandler
  -> [U.Elt U.ExpandedOp]
  -> U.Value
  -> Sing kt
  -> TypeCheckInstr [(T.Value kt, T.Value vt)]
typeCheckMapVal tcDo mels sq kt = withComparable kt sq kt $ do
  instrPos <- ask
  ks <- typeCheckValsImpl @kt tcDo (map (\(U.Elt k _) -> k) mels)
  vals <- typeCheckValsImpl @vt tcDo (map (\(U.Elt _ v) -> v) mels)
  ksS <- liftEither $ ensureDistinctAsc id ks
        `onLeft` \msg -> TCFailedOnValue sq (fromSingT kt) msg instrPos Nothing
  pure $ zip ksS vals

typeCheckValsImpl
  :: forall t . (SingI t)
  => TcInstrHandler
  -> [U.Value]
  -> TypeCheckInstr [T.Value t]
typeCheckValsImpl tcDo mvs =
  fmap reverse $ foldM check [] mvs
  where
    check :: [T.Value t] -> U.Value -> TypeCheckInstr [T.Value t]
    check res mv = do
      v <- typeCheckValImpl @t tcDo mv
      pure (v : res)