-- | Parsing of Michelson types.

module Michelson.Parser.Type
  ( type_
  , comparable
  ) where

import Prelude hiding (many, note, some, try)

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

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

type_ :: Parser Type
type_ = snd <$> typeInner (pure noAnn)

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

t_operator :: Parser FieldAnn -> Parser (FieldAnn, Type)
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', Type ty' _) = mkGenericTree (mergeTwo isOr) (ty :| tys)
      f'' <- mergeAnnots f f'
      return (f'', Type ty' t)
    Just (res, _) -> do
      return res
    Nothing -> do
      return (f, Type TUnit t)
  where
    mergeTwo isOr _ (l, a) (r, b) =
      (noAnn, Type ((if isOr then TOr else TPair) l r a b) noAnn)

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

typeInner
  :: Parser FieldAnn -> Parser (FieldAnn, Type)
typeInner fp = lexeme $ choice $ (\x -> x fp) <$>
  [ t_ct, t_key, t_unit, t_signature, t_chain_id
  , t_option, t_list, t_set
  , t_operation, t_contract, t_pair, t_or
  , t_lambda, t_map, t_big_map, t_view
  , t_void, t_letType
  , t_operator
  , const (customFailure UnknownTypeException)
  ]

----------------------------------------------------------------------------
-- Comparable types
----------------------------------------------------------------------------

comparable :: Parser Comparable
comparable = mparens $ Comparable <$> ct <*> noteDef

t_ct :: (Default a) => Parser a -> Parser (a, Type)
t_ct fp = do ct' <- ct; (f,t) <- fieldType fp; return (f, Type (Tc ct') t)

ct :: Parser CT
ct =  word' "Int" CInt
  <|> word' "Nat" CNat
  <|> word' "String" CString
  <|> word' "Bytes" CBytes
  <|> word' "Mutez" CMutez
  <|> word' "Bool" CBool
  <|> word' "KeyHash" CKeyHash <|> word "key_hash" CKeyHash
  <|> word' "Timestamp" CTimestamp
  <|> word' "Address" CAddress

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

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

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

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

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

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

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

t_unit :: (Default a) => Parser a -> Parser (a, Type)
t_unit fp = do
  symbol' "Unit" <|> symbol' "()"
  (f,t) <- fieldType fp
  return (f, Type TUnit t)

t_pair_like
  :: (Default a)
  => (FieldAnn -> FieldAnn -> Type -> Type -> T)
  -> Parser a
  -> Parser (a, Type)
t_pair_like mkPair fp = do
  (f, t) <- fieldType fp
  (l, a) <- field
  (r, b) <- field
  return (f, Type (mkPair l r a b) t)

t_pair :: (Default a) => Parser a -> Parser (a, Type)
t_pair fp = do
  symbol' "Pair"
  t_pair_like TPair fp

t_or :: (Default a) => Parser a -> Parser (a, Type)
t_or fp = do
  symbol' "Or"
  t_pair_like TOr fp

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

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

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

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

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

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

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

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

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

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

t_letType :: Default fp => Parser fp -> Parser (fp, Type)
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)