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(..), 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, CValue(..), EpAddress(..), Notes(..), Sing(..), Value'(..), 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
  :: forall op ct
  .  U.Value' op -> Sing ct -> Either (U.Value' op, TCTypeError) (CValue ct)
typeCheckCValue val ct = case (val, ct) of
  (U.ValueInt i, SCInt) -> pure $ CvInt i
  (U.ValueInt i, SCNat)
    | i >= 0 -> pure $ CvNat (fromInteger i)
    | otherwise -> Left (U.ValueInt i, NegativeNat)
  (v@(U.ValueInt i), SCMutez) -> do
    mtz <- maybeToRight (v, InvalidTimestamp) . mkMutez $ fromInteger i
    pure $ CvMutez mtz
  (U.ValueString s, SCString) ->
    pure $ CvString s
  (v@(U.ValueString s), SCAddress) -> do
    addr <- T.parseEpAddress (unMText s) `onLeft` ((v, ) . InvalidAddress)
    pure $ CvAddress addr
  (v@(U.ValueString s), SCKeyHash) -> do
    kHash <- parseKeyHash (unMText s) `onLeft` ((v, ) . InvalidKeyHash)
    pure $ CvKeyHash kHash
  (v@(U.ValueString s), SCTimestamp) -> do
    tstamp <- maybeToRight (v, InvalidTimestamp) . parseTimestamp $ unMText s
    pure $ CvTimestamp tstamp
  (U.ValueInt i, SCTimestamp) ->
    pure $ CvTimestamp (timestampFromSeconds i)
  (U.ValueBytes (U.InternalByteString s), SCBytes) ->
    pure $ CvBytes s
  (U.ValueTrue, SCBool) -> pure $ CvBool True
  (U.ValueFalse, SCBool) -> pure $ CvBool False
  (v, t) ->
    Left $ (v, InvalidValueType $ T.Tc $ fromSingCT t)

typeCheckCVals
  :: forall t op
  .  [U.Value' op]
  -> Sing t
  -> Either (U.Value' op, TCTypeError) [CValue t]
typeCheckCVals mvs t = traverse (`typeCheckCValue` t) mvs

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
  (mv, t@(STc ct)) -> do
    case typeCheckCValue mv ct of
      Left (uval, err) -> tcFailedOnValue uval (fromSingT t) "" (Just err)
      Right v -> pure $ VC v
  (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 :: 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 @'T.TUnit
        pure $ VContract addr T.sepcPrimitive
      ContractAddress ca ->
        case M.lookup ca contracts of
          -- Wrapping into 'ParamNotesUnsafe' is safe because originated contract has
          -- valid parameter type
          Just (AsUTypeExt 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 cpNotes 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 _) ->
    pure $ T.VSet S.empty

  (sq@(U.ValueSeq (toList -> mels)), STSet vt) -> do
    instrPos <- ask
    els <- liftEither $ typeCheckCVals mels 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

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

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

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

  (sq@(U.ValueMap (toList -> mels)), STBigMap kt _) -> 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 <- typeCheckImpl tcDo mp ((starNotes @it, def) ::& SNil)
    case instr of
      lam ::: (lo :: HST lo) -> 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

-- | 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)
  => TcInstrHandler
  -> [U.Elt U.ExpandedOp]
  -> U.Value
  -> Sing kt
  -> TypeCheckInstr [(CValue kt, T.Value vt)]
typeCheckMapVal tcDo mels sq kt = do
  instrPos <- ask
  ks <- liftEither $ typeCheckCVals (map (\(U.Elt k _) -> k) mels) 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)
  ksS <- liftEither $ ensureDistinctAsc id ks
        `onLeft` \msg -> TCFailedOnValue sq (fromSingT $ STc 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)