-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Test.Util.Annotation
  ( FieldAnnTree (..)
  , TypeAnnTree (..)
  , extractFieldAnnTree
  , extractTypeAnnTree
  ) where

import Morley.Michelson.Typed
import Morley.Michelson.Untyped (FieldAnn, TypeAnn, noAnn)

-- | Helper datatype which contains field annotations from the given type
data FieldAnnTree t where
  FALeaf :: FieldAnnTree t
  FANodeOr :: FieldAnn -> FieldAnnTree a
           -> FieldAnn -> FieldAnnTree b
           -> FieldAnnTree ('TOr a b)
  FANodePair :: FieldAnn -> FieldAnnTree a
             -> FieldAnn -> FieldAnnTree b
             -> FieldAnnTree ('TPair a b)

deriving stock instance Eq (FieldAnnTree t)
deriving stock instance Show (FieldAnnTree t)

extractFieldAnnTree :: Notes t -> FieldAnnTree t
extractFieldAnnTree = \case
  NTOr _ lann rann lnotes rnotes ->
    FANodeOr lann (extractFieldAnnTree lnotes) rann (extractFieldAnnTree rnotes)
  NTPair _ lann rann _ _ lnotes rnotes
    | lann == noAnn && rann == noAnn -> FALeaf
    | otherwise -> FANodePair lann (extractFieldAnnTree lnotes) rann (extractFieldAnnTree rnotes)
  _ ->
    FALeaf

-- | Helper datatype which contains type annotations from the given type
data TypeAnnTree t where
  TALeaf :: TypeAnn -> TypeAnnTree t
  TANodeOption :: TypeAnn -> TypeAnnTree a -> TypeAnnTree ('TOption a)
  TANodePair :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TPair a b)
  TANodeOr :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TOr a b)
  TANodeLambda :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TLambda a b)
  TANodeList :: TypeAnn -> TypeAnnTree a -> TypeAnnTree ('TList a)
  TANodeMap :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TMap a b)
  TANodeBigMap :: TypeAnn ->TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TBigMap a b)

deriving stock instance Eq (TypeAnnTree t)
deriving stock instance Show (TypeAnnTree t)

extractTypeAnnTree :: Notes t -> TypeAnnTree t
extractTypeAnnTree = \case
  NTKey ta -> TALeaf ta
  NTUnit ta -> TALeaf ta
  NTSignature ta -> TALeaf ta
  NTOption ta n1 -> TANodeOption ta (extractTypeAnnTree n1)
  NTList ta n1 -> TANodeList ta (extractTypeAnnTree n1)
  NTSet ta _ -> TALeaf ta
  NTOperation ta -> TALeaf ta
  NTContract ta _ -> TALeaf ta
  NTTicket ta _ -> TALeaf ta
  NTPair ta _ _ _ _ n1 n2 -> TANodePair ta (extractTypeAnnTree n1) (extractTypeAnnTree n2)
  NTOr ta _ _ n1 n2 -> TANodeOr ta (extractTypeAnnTree n1) (extractTypeAnnTree n2)
  NTLambda ta n1 n2 -> TANodeLambda ta (extractTypeAnnTree n1) (extractTypeAnnTree n2)
  NTMap ta n1 n2 -> TANodeMap ta (extractTypeAnnTree n1) (extractTypeAnnTree n2)
  NTBigMap ta n1 n2 -> TANodeBigMap ta (extractTypeAnnTree n1) (extractTypeAnnTree n2)
  NTChainId ta -> TALeaf ta
  NTInt ta -> TALeaf ta
  NTNat ta -> TALeaf ta
  NTString ta -> TALeaf ta
  NTBytes ta -> TALeaf ta
  NTMutez ta -> TALeaf ta
  NTBool ta -> TALeaf ta
  NTKeyHash ta -> TALeaf ta
  NTBls12381Fr ta -> TALeaf ta
  NTBls12381G1 ta -> TALeaf ta
  NTBls12381G2 ta -> TALeaf ta
  NTTimestamp ta -> TALeaf ta
  NTAddress ta -> TALeaf ta
  NTNever ta -> TALeaf ta
  NTChest ta -> TALeaf ta
  NTChestKey ta -> TALeaf ta
  NTSaplingState ta _ -> TALeaf ta
  NTSaplingTransaction ta _ -> TALeaf ta