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

-- | Parsing of Michelson types.

module Morley.Michelson.Parser.Type
  ( type_
  , field
  ) where

import Prelude hiding (note, some, try)

import Data.Default (Default, def)
import qualified Data.Map as Map
import Fmt (pretty)
import Text.Megaparsec (choice, customFailure, label, sepBy)

import Morley.Michelson.Let (LetType(..))
import Morley.Michelson.Parser.Annotations
import Morley.Michelson.Parser.Error
import Morley.Michelson.Parser.Helpers
import Morley.Michelson.Parser.Lexer
import Morley.Michelson.Parser.Types (Parser, letTypes)
import Morley.Michelson.Untyped
import Morley.Util.Generic

-- | This parses arbitrary type expressions.
--
-- Note that this includes parenthesized ones for efficiency, see 't_operator'.
-- That is to say, @int@, @(int)@, @((int))@, etc will match with this parser and produce @TInt@.
type_ :: Parser Ty
type_ = snd <$> typeInner (pure noAnn)

field :: Parser (FieldAnn, Ty)
field = typeInner note

-- | 't_operator' parses tuples @(a, b, c)@, variants @(a | b | c)@, and also type expressions in
-- parentheses @(a)@ and unit @()@. This is done this way for performance considerations.
--
-- Consequently, 't_unit' doesn't bother with parsing @()@.
t_operator :: Parser FieldAnn -> Parser (FieldAnn, Ty)
t_operator fp = do
  whole <- parens do
    optional do
      ty   <- field
      rest <- optional do
        isOr <- (symbol' "|" >> return True)
            <|> (symbol' "," >> return False)
        others <- field `sepBy` symbol' if isOr then "|" else ","
        return (isOr, others)
      return (ty, rest)

  (f, t) <- fieldType fp
  case whole of
    Just (ty, Just (isOr, tys)) -> do
      let (f', Ty ty' _) = mkGenericTree (mergeTwo isOr) (ty :| tys)
      f'' <- mergeAnnots f f'
      return (f'', Ty ty' t)
    Just (res, _) -> do
      return res
    Nothing -> do
      return (f, Ty TUnit t)
  where
    mergeTwo isOr _ (l, a) (r, b) =
      (noAnn, Ty ((if isOr then TOr l r else TPair l r noAnn noAnn) a b) noAnn)

    mergeAnnots l r
      | l == def  = return r
      | r == def  = return l
      | otherwise = customFailure ExcessFieldAnnotation

typeInner
  :: Parser FieldAnn -> Parser (FieldAnn, Ty)
typeInner fp = label "type" $ lexeme $ choice $ (\x -> x fp) <$>
  [ t_operator
  , t_int, t_nat, t_string, t_bytes, t_mutez, t_bool
  , t_keyhash, t_timestamp, t_address
  , t_key, t_unit, t_never, t_signature, t_chain_id
  , t_bls12381fr, t_bls12381g1, t_bls12381g2
  , t_option, t_list, t_set
  , t_operation, t_contract, t_ticket, t_pair, t_or
  , t_lambda, t_map, t_big_map, t_view
  , t_void, t_letType
  ]

----------------------------------------------------------------------------
-- Non-comparable types
----------------------------------------------------------------------------

mkType :: T -> (a, TypeAnn) -> (a, Ty)
mkType t (a, ta) = (a, Ty t ta)


t_int :: (Default a) => Parser a -> Parser (a, Ty)
t_int fp = word' "Int" (mkType TInt) <*> fieldType fp

t_nat :: (Default a) => Parser a -> Parser (a, Ty)
t_nat fp = word' "Nat" (mkType TNat) <*> fieldType fp

t_string :: (Default a) => Parser a -> Parser (a, Ty)
t_string fp = word' "String" (mkType TString) <*> fieldType fp

t_bytes :: (Default a) => Parser a -> Parser (a, Ty)
t_bytes fp = word' "Bytes" (mkType TBytes) <*> fieldType fp

t_mutez :: (Default a) => Parser a -> Parser (a, Ty)
t_mutez fp = word' "Mutez" (mkType TMutez) <*> fieldType fp

t_bool :: (Default a) => Parser a -> Parser (a, Ty)
t_bool fp = word' "Bool" (mkType TBool) <*> fieldType fp

t_keyhash :: (Default a) => Parser a -> Parser (a, Ty)
t_keyhash fp = ((word' "KeyHash" (mkType TKeyHash)) <|> (word "key_hash" (mkType TKeyHash))) <*> fieldType fp

t_timestamp :: (Default a) => Parser a -> Parser (a, Ty)
t_timestamp fp = word' "Timestamp" (mkType TTimestamp) <*> fieldType fp

t_address :: (Default a) => Parser a -> Parser (a, Ty)
t_address fp = word' "Address" (mkType TAddress) <*> fieldType fp

t_key :: (Default a) => Parser a -> Parser (a, Ty)
t_key fp = word' "Key" (mkType TKey) <*> fieldType fp

t_signature :: (Default a) => Parser a -> Parser (a, Ty)
t_signature fp = word' "Signature" (mkType TSignature) <*> fieldType fp

t_bls12381fr :: (Default a) => Parser a -> Parser (a, Ty)
t_bls12381fr fp = do
  symbol' "bls12_381_fr" <|> symbol' "Bls12381Fr"
  mkType TBls12381Fr <$> fieldType fp

t_bls12381g1 :: (Default a) => Parser a -> Parser (a, Ty)
t_bls12381g1 fp = do
  symbol' "bls12_381_g1" <|> symbol' "Bls12381G1"
  mkType TBls12381G1 <$> fieldType fp

t_bls12381g2 :: (Default a) => Parser a -> Parser (a, Ty)
t_bls12381g2 fp = do
  symbol' "bls12_381_g2" <|> symbol' "Bls12381G2"
  mkType TBls12381G2 <$> fieldType fp

t_chain_id :: (Default a) => Parser a -> Parser (a, Ty)
t_chain_id fp = do
  symbol' "ChainId" <|> symbol' "chain_id"
  mkType TChainId <$> fieldType fp

t_operation :: (Default a) => Parser a -> Parser (a, Ty)
t_operation fp = word' "Operation" (mkType TOperation) <*> fieldType fp

t_contract :: (Default a) => Parser a -> Parser (a, Ty)
t_contract fp = do
  symbol' "Contract"
  (f, t) <- fieldType fp
  a <- type_
  return (f, Ty (TContract a) t)

t_ticket :: (Default a) => Parser a -> Parser (a, Ty)
t_ticket fp = do
  symbol' "Ticket"
  (f, t) <- fieldType fp
  a <- type_
  return (f, Ty (TTicket a) t)

-- | Parses a @unit@ type. Unit type admits two variants of syntax, @unit@ and @()@. This parser
-- handles only the former, the latter is handled in 't_operator'
t_unit :: (Default a) => Parser a -> Parser (a, Ty)
t_unit fp = do
  symbol' "Unit"
  (f,t) <- fieldType fp
  return (f, Ty TUnit t)

t_never :: (Default a) => Parser a -> Parser (a, Ty)
t_never fp = do
  symbol' "Never" <|> symbol' "⊥"
  (f,t) <- fieldType fp
  return (f, Ty TNever t)

t_pair :: (Default a) => Parser a -> Parser (a, Ty)
t_pair fp = do
  symbol' "Pair"
  (fieldAnn, typeAnn) <- fieldType fp
  fields <- many field
  tPair <- go fields
  pure $ (fieldAnn, Ty tPair typeAnn)
  where
    go :: [(FieldAnn, Ty)] -> Parser T
    go = \case
      [] -> fail "The 'pair' type expects at least 2 type arguments, but 0 were given."
      [(_, t)] -> fail $ "The 'pair' type expects at least 2 type arguments, but only 1 was given: '" <> pretty t <> "'."
      [(fieldAnnL, typeL), (fieldAnnR, typeR)] ->
        pure $ TPair fieldAnnL fieldAnnR noAnn noAnn typeL typeR
      (fieldAnnL, typeL) : fields -> do
        rightCombedT <- go fields
        pure $ TPair fieldAnnL noAnn noAnn noAnn typeL (Ty rightCombedT noAnn)

t_or :: (Default a) => Parser a -> Parser (a, Ty)
t_or fp = do
  symbol' "Or"
  (f, t) <- fieldType fp
  (l, a) <- field
  (r, b) <- field
  return (f, Ty (TOr l r a b) t)

t_option :: (Default a) => Parser a -> Parser (a, Ty)
t_option fp = do
  symbol' "Option"
  (f, t) <- fieldType fp
  a <- snd <$> typeInner (pure noAnn)
  return (f, Ty (TOption a) t)

t_lambda :: (Default a) => Parser a -> Parser (a, Ty)
t_lambda fp = core <|> slashLambda
  where
    core = do
      symbol' "Lambda"
      (f, t) <- fieldType fp
      a <- type_
      b <- type_
      return (f, Ty (TLambda a b) t)
    slashLambda = do
      symbol "\\"
      (f, t) <- fieldType fp
      a <- type_
      symbol "->"
      b <- type_
      return (f, Ty (TLambda a b) t)

-- Container types
t_list :: (Default a) => Parser a -> Parser (a, Ty)
t_list fp = core <|> bracketList
  where
    core = do
      symbol' "List"
      (f, t) <- fieldType fp
      a <- type_
      return (f, Ty (TList a) t)
    bracketList = do
      a <- brackets type_
      (f, t) <- fieldType fp
      return (f, Ty (TList a) t)

t_set :: (Default a) => Parser a -> Parser (a, Ty)
t_set fp = core <|> braceSet
  where
    core = do
      symbol' "Set"
      (f, t) <- fieldType fp
      a <- type_
      return (f, Ty (TSet a) t)
    braceSet = do
      a <- braces type_
      (f, t) <- fieldType fp
      return (f, Ty (TSet a) t)

t_map_like
  :: Default a
  => Parser a -> Parser (Ty, Ty, a, TypeAnn)
t_map_like fp = do
  (f, t) <- fieldType fp
  a <- type_
  b <- type_
  return (a, b, f, t)

t_map :: (Default a) => Parser a -> Parser (a, Ty)
t_map fp = do
  symbol' "Map"
  (a, b, f, t) <- t_map_like fp
  return (f, Ty (TMap a b) t)

t_big_map :: (Default a) => Parser a -> Parser (a, Ty)
t_big_map fp = do
  symbol' "BigMap" <|> symbol "big_map"
  (a, b, f, t) <- t_map_like fp
  return (f, Ty (TBigMap a b) t)

----------------------------------------------------------------------------
-- Non-standard types (Morley extensions)
----------------------------------------------------------------------------

t_view :: Default a => Parser a -> Parser (a, Ty)
t_view fp = do
  symbol' "View"
  a <- type_
  r <- type_
  (f, t) <- fieldType fp
  let c' = Ty (TContract r) noAnn
  return (f, Ty (TPair noAnn noAnn noAnn noAnn a c') t)

t_void :: Default a => Parser a -> Parser (a, Ty)
t_void fp = do
  symbol' "Void"
  a <- type_
  b <- type_
  (f, t) <- fieldType fp
  let c = Ty (TLambda b b) noAnn
  return (f, Ty (TPair noAnn noAnn noAnn noAnn a c) t)

t_letType :: Default fp => Parser fp -> Parser (fp, Ty)
t_letType fp = do
  lts <- asks letTypes
  lt <- ltSig <$> (mkLetType lts)
  f <- parseDef fp
  return (f, lt)

mkLetType :: Map Text LetType -> Parser LetType
mkLetType lts = choice $ mkParser ltName <$> (Map.elems lts)