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

module Morley.Michelson.Typed.AnnotatedValue
  ( AnnotatedValue (..)
  , SomeAnnotatedValue(..)
  , getT
  -- * Optics
  , value
  , castTo
  , variant
  , fields
  , field
  , nodes
  , node
  , asList
  , asMap
  ) where

import Control.Lens
  (Fold, IndexedFold, IndexedTraversal', Prism, folding, ifolding, index, indexed, lens, prism,
  (:~:)(Refl))
import Data.Map qualified as Map
import Data.Singletons (SingI, demote)
import Fmt (Buildable(..), (+|), (|+))

import Morley.Michelson.Typed.Aliases (Value)
import Morley.Michelson.Typed.Annotation (Notes(..), mkUType)
import Morley.Michelson.Typed.Convert ()
import Morley.Michelson.Typed.Haskell.Value (IsoValue, ToT, fromVal, toVal)
import Morley.Michelson.Typed.T (T)
import Morley.Michelson.Typed.Value
import Morley.Michelson.Untyped (FieldAnn, noAnn)
import Morley.Util.Sing (eqI)

-- $setup
-- >>> import qualified Data.Map as Map
-- >>> import Morley.Michelson.Parser (notes)
-- >>> import Morley.Michelson.Untyped (fieldAnnQ)
-- >>> :m +Morley.Michelson.Typed Morley.Michelson.Text Control.Lens Fmt

data AnnotatedValue v = AnnotatedValue
  { avNotes :: Notes (ToT v)
  , avValue :: Value (ToT v)
  }
  deriving stock (Eq, Show)

instance Buildable (AnnotatedValue v) where
  build (AnnotatedValue n v) = build v |+ " :: " +| build (mkUType n)

data SomeAnnotatedValue where
  SomeAnnotatedValue :: forall t. SingI t => Notes t -> Value t -> SomeAnnotatedValue

deriving stock instance Show SomeAnnotatedValue

instance Eq SomeAnnotatedValue where
  SomeAnnotatedValue n1 (v1 :: Value t1) == SomeAnnotatedValue n2 (v2 :: Value t2) =
    case eqI @t1 @t2 of
      Nothing -> False
      Just Refl -> n1 == n2 && v1 == v2

instance Buildable SomeAnnotatedValue where
  build (SomeAnnotatedValue n v) = build v |+ " :: " +| build (mkUType n)

getT :: SomeAnnotatedValue -> T
getT (SomeAnnotatedValue _ (_ :: Value t)) = demote @t

----------------------------------------------------------------------------
-- Optics
----------------------------------------------------------------------------

-- | Extract the value out of an annotated value.
value :: IsoValue v => Lens' (AnnotatedValue v) v
value =
  lens
    do fromVal . avValue
    do \av v -> av { avValue = toVal v }

-- | Attempts to cast to the given type.
castTo
  :: forall v1 v2. (IsoValue v1, IsoValue v2)
  => Prism SomeAnnotatedValue SomeAnnotatedValue (AnnotatedValue v1) (AnnotatedValue v2)
castTo =
  prism
    do \(AnnotatedValue n v) -> SomeAnnotatedValue n v
    do \sav@(SomeAnnotatedValue n (v :: Value t)) ->
        case eqI @t @(ToT v1) of
          Just Refl -> Right $ AnnotatedValue n v
          Nothing -> Left sav

{- | Matches on a (possibly nested) @or@, and focuses the value if its
field annotation is a match.

Note that field annotations do not have to be unique.
If two @or@ branches have the same field annotation, `variant` will match on either branch.

>>> :{
val = SomeAnnotatedValue
  [notes|or (int %first) (or (nat %second) (int %third))|]
  (toVal (Right (Left 3) :: Either Integer (Either Natural Integer)))
:}

>>> val ^? variant [fieldAnnQ|second|] . to pretty
Just "3 :: nat"

>>> val ^? variant [fieldAnnQ|first|] . to pretty
Nothing
-}
variant :: FieldAnn -> Traversal' SomeAnnotatedValue SomeAnnotatedValue
variant fieldAnn (handler :: SomeAnnotatedValue -> (f SomeAnnotatedValue)) = \case
  SomeAnnotatedValue n0 v0 -> go noAnn (n0, v0)
  where
    go :: FieldAnn -> (Notes t, Value t) -> f SomeAnnotatedValue
    go currentFieldAnn (n, v)
      -- If this node does not have a field ann, then recurse.
      | currentFieldAnn == noAnn =
          case (n, v) of
            (NTOr ta fal far nl nr, VOr (Left vl)) -> do
              go fal (nl, vl) <&> \(SomeAnnotatedValue nl' vl') ->
                SomeAnnotatedValue (NTOr ta fal far nl' nr) (VOr (Left vl'))
            (NTOr ta fal far nl nr, VOr (Right vr)) -> do
              go far (nr, vr) <&> \(SomeAnnotatedValue nr' vr') ->
                SomeAnnotatedValue (NTOr ta fal far nl nr') (VOr (Right vr'))
            _ ->
              -- We've reached a leaf.
              pure $ withValueTypeSanity v $ SomeAnnotatedValue n v

      -- We've encountered a node labelled with a field ann, and it's a match.
      | currentFieldAnn == fieldAnn = handler $ withValueTypeSanity v $ SomeAnnotatedValue n v
      -- We've encountered a node labelled with a field ann, but it's not a match.
      | otherwise = pure $ withValueTypeSanity v $ SomeAnnotatedValue n v

{- | Traverses all nodes in a tree of pairs that have a field annotation.
The elements are indexed by their field annotations.

Note that sub-trees are not inspected.
For example, for this type:

> pair
>   (int %storageField1)
>   (pair %storageField2
>     (int %nestedField1)
>     (string %nestedField2)
>   )

'fields' will:

* ignore the top-level node, because it does not have a field annotation.
* traverse the @storageField1@ and @storageField2@ fields.
* ignore @nestedField1@ and @nestedField2@, because they're in a sub-tree of a traversed node.

Note that field annotations do not have to be unique.
If two fields have the same field annotation, both will be indexed/traversed.

>>> :{
val = SomeAnnotatedValue
  [notes|pair (int %first) (nat %second) (int %third)|]
  (toVal ((1, 2, 3) :: (Integer, Natural, Integer)))
:}

>>> mapM_ print $ map (bimap pretty pretty) $ val ^@.. fields
("%first","1 :: int")
("%second","2 :: nat")
("%third","3 :: int")
-}
fields :: IndexedTraversal' FieldAnn SomeAnnotatedValue SomeAnnotatedValue
fields (handler :: p SomeAnnotatedValue (f SomeAnnotatedValue)) = \case
  SomeAnnotatedValue n0 v0 -> go noAnn (n0, v0)
  where
    go :: FieldAnn -> (Notes t, Value t) -> f SomeAnnotatedValue
    go currentFieldAnn (n, v)
      | currentFieldAnn == noAnn =
          case (n, v) of
            (NTPair ta fal far val var nl nr, VPair (vl, vr)) -> do
              savL <- go fal (nl, vl)
              savR <- go far (nr, vr)
              pure $
                case (savL, savR) of
                  (SomeAnnotatedValue nl' vl', SomeAnnotatedValue nr' vr') ->
                    SomeAnnotatedValue (NTPair ta fal far val var nl' nr') (VPair (vl', vr'))
            _ ->
              -- We've reached a leaf.
              pure $ withValueTypeSanity v $ SomeAnnotatedValue n v
      | otherwise =
          indexed handler currentFieldAnn $ withValueTypeSanity v $ SomeAnnotatedValue n v

{- | Looks up a member of a tree of pairs by its field annotation.
Note that sub-trees of nodes with field annotations are not inspected.

Note that field annotations do not have to be unique.
If two fields have the same field annotation, both will be indexed/traversed.

>>> :{
val = SomeAnnotatedValue
  [notes|pair (int %first) (nat %second) (int %third)|]
  (toVal ((1, 2, 3) :: (Integer, Natural, Integer)))
:}

>>> val ^? field [fieldAnnQ|second|] . castTo @Natural . value
Just 2
-}
field :: FieldAnn -> Traversal' SomeAnnotatedValue SomeAnnotatedValue
field fieldAnn = fields . index fieldAnn

{- | Lists all nodes in a tree of pairs.
The elements are indexed by their field annotations.

Note that field annotations do not have to be unique.
If two fields have the same field annotation, both will be indexed.

>>> :{
val = SomeAnnotatedValue
  [notes|pair (int %first) (nat %second) (int %third)|]
  (toVal ((1, 2, 3) :: (Integer, Natural, Integer)))
:}

>>> mapM_ print $ map (bimap pretty pretty) $ val ^@.. nodes
("%first","1 :: int")
("%","Pair 2 3 :: pair (nat %second) (int %third)")
("%second","2 :: nat")
("%third","3 :: int")
-}
nodes :: IndexedFold FieldAnn SomeAnnotatedValue SomeAnnotatedValue
nodes = ifolding $ \case
  SomeAnnotatedValue n0 v0 -> go (n0, v0)
  where
    go :: (Notes t, Value t) -> [(FieldAnn, SomeAnnotatedValue)]
    go (n, v) = case (n, v) of
      (NTPair _ fal far _ _ nl nr, VPair (vl, vr)) ->
        (fal, withValueTypeSanity vl $ SomeAnnotatedValue nl vl) :
        (far, withValueTypeSanity vr $ SomeAnnotatedValue nr vr) :
        go (nl, vl) <>
        go (nr, vr)
      _ -> []

-- | Looks up nodes in a tree of pairs with the given field annotation.
node :: FieldAnn -> Fold SomeAnnotatedValue SomeAnnotatedValue
node fieldAnn = nodes . index fieldAnn

-- | Casts a 'SomeAnnotatedValue' to a list.
--
-- >>> list = [ [mt|hello|], [mt|world|] ]
-- >>> val = SomeAnnotatedValue starNotes (toVal list)
-- >>> val ^? asList . ix 1 . castTo @MText . value . to pretty
-- Just "world"
asList :: Fold SomeAnnotatedValue [SomeAnnotatedValue]
asList = folding @Maybe \case
  SomeAnnotatedValue (NTList _ xNotes) (VList xs) ->
    Just $ xs <&> \x -> SomeAnnotatedValue xNotes x
  _ -> Nothing

-- | Casts a 'SomeAnnotatedValue' to a map.
--
-- >>> map = Map.fromList @Integer @Bool [(10, False), (20, True)]
-- >>> val = SomeAnnotatedValue starNotes (toVal map)
-- >>> val ^? asMap @Integer . ix 20 . castTo @Bool . value
-- Just True
asMap
  :: forall key. (IsoValue key, Ord key)
  => Fold SomeAnnotatedValue (Map key SomeAnnotatedValue)
asMap = folding @Maybe \case
  SomeAnnotatedValue (NTMap _ _ vNotes) (VMap vmap) -> convert vNotes vmap
  SomeAnnotatedValue (NTBigMap _ _ vNotes) (VBigMap _ vmap) -> convert vNotes vmap
  _ -> Nothing
  where
    convert
      :: forall k v. (SingI k, SingI v)
      => Notes v
      -> Map (Value k) (Value v)
      -> Maybe (Map key SomeAnnotatedValue)
    convert vNotes vmap =
      eqI @(ToT key) @k <&> \Refl ->
        vmap
          & Map.mapKeys (\k -> fromVal @key k)
          <&> \v -> SomeAnnotatedValue vNotes v