{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | This module implements a search in a typelevel tree
--   and offers a handy interface via @-XOverloaedRecordDot@ and 'GHC.Records.HasField'.
--   We can search in the tree via BFS or DFS.
--   Performance wise this doesn't make a difference, as the search is performed at
--   compile time anyway however, it can change the semantics if the tree contains an
--   element more than once, see the example in 'getElem'.
module Data.HTree.Labeled
  ( -- * Interface

    -- ** Labeled types
    Labeled (..)
  , pattern HNodeL
  , pattern HLeafL
  , type TyNodeL

    -- ** Getting elements
  , getElem
  , getElem'
  , SearchStrategy (..)

    -- ** Reexports
  , HasField (..)
  , Proxy (..)

    -- * Internal
  , HasField' (..)
  , Decide (..)
  , Elem
  , AnyElem
  , getElemWithPath
  )
where

import Data.HTree.Families (Not, type (||))
import Data.HTree.List (HList (HCons, HNil))
import Data.HTree.Tree (HForest, HTree (HNode), Path (Deeper, Farther, Here), TyForest, TyTree (TyNode))
import Data.Kind (Constraint, Type)
import Data.Proxy (Proxy (Proxy))
import GHC.Generics (Generic)
import GHC.Records (HasField (getField))

-- | a type syonym that allows for easy construction of TyTrees that have labeled nodes
type TyNodeL l a = 'TyNode (Labeled l a)

-- | a pattern that allows for direct construction and destruction of nodes with
--   labels
pattern HNodeL :: forall l a f ts. Functor f => f a -> HForest f ts -> HTree f ('TyNode (Labeled l a) ts)
pattern HNodeL x ts <- (HNode (fmap unLabel -> x) ts)
  where
    HNodeL x ts = HNode (MkLabeled <$> x) ts

{-# COMPLETE HNodeL #-}

-- | a labeled HNode Leaf
pattern HLeafL :: forall l a f. Functor f => f a -> HTree f ('TyNode (Labeled l a) '[])
pattern HLeafL x <- (HNode (fmap unLabel -> x) HNil)
  where
    HLeafL x = HNode (MkLabeled <$> x) HNil

-- | a newtype that is labeled with some typelevel tag
type Labeled :: forall l. l -> Type -> Type
newtype Labeled l a = MkLabeled {unLabel :: a}
  deriving stock (Show, Eq, Ord, Functor, Foldable, Traversable, Generic)

-- | gets an element given a path into the tree
getElemWithPath :: forall typ t f. Path typ t -> HTree f t -> f typ
getElemWithPath Here (HNode e _) = e
getElemWithPath (Farther pt) (HNode e (_ `HCons` ts)) = getElemWithPath pt (HNode e ts)
getElemWithPath (Deeper pt) (HNode _ (t `HCons` _)) = getElemWithPath pt t

-- | searches a tree for an element and returns that element
getElem' :: forall {proxy} strat typ t f. HasField' (strat :: SearchStrategy) typ t => proxy strat -> HTree f t -> f typ
getElem' _ = getElemWithPath (evidence @strat @typ @t Proxy)

-- | searches a tree for an element and returns that element, specialised to 'Labeled' and unwraps
--
-- >>> import Data.Functor.Identity
-- >>> type T = TyNodeL "top" Int [ TyNodeL "inter" Int '[ TyNodeL "foo" Int '[]], TyNodeL "foo" Int '[]]
-- >>> t :: HTree Identity T = 42 `HNodeL` HNodeL 4 (HNodeL 69 HNil `HCons` HNil) `HCons` HNodeL 67 HNil `HCons` HNil
-- >>> getElem @'DFS @"foo" @Int Proxy t
-- Identity 69
-- >>> getElem @'BFS @"foo" @Int Proxy t
-- Identity 67
getElem
  :: forall {proxy} strat l typ t f
   . ( HasField' (strat :: SearchStrategy) (Labeled l typ) t
     , Functor f
     )
  => proxy strat
  -> HTree f t
  -> f typ
getElem _ = fmap unLabel . getElemWithPath (evidence @strat @(Labeled l typ) @t Proxy)

-- the default behaviour is a breadth first search
instance (HasField' 'BFS (Labeled l typ) t, Functor f) => HasField l (HTree f t) (f typ) where
  getField = getElem @'BFS @l @typ @t Proxy

-- | simple typelevel predicate that tests whether some element is in a tree
type Elem :: forall k. k -> TyTree k -> Bool
type family Elem typ t where
  Elem a ('TyNode a ts) = 'True
  Elem a ('TyNode a' ts) = AnyElem a ts
  Elem a t = 'False

-- | typelevel predicate that tests whether the element is in any of the
--   subtrees
type AnyElem :: forall k. k -> TyForest k -> Bool
type family AnyElem typ ts where
  AnyElem a (t : ts) = Elem a t || AnyElem a ts
  AnyElem a '[] = 'False

-- | the search strategy used in 'HasField'', this is intended to be used only as a DataKind
type SearchStrategy :: Type
data SearchStrategy = DFS | BFS

-- | This is the helper class that creates evidence,
--   it implements a DFS together with Decide
type HasField' :: SearchStrategy -> Type -> TyTree Type -> Constraint
class HasField' strat typ t | strat t -> typ where
  evidence :: forall {proxy}. proxy strat -> Path typ t

-- | Together with HasField' implements a DFS in the tree
type Decide :: SearchStrategy -> Bool -> Type -> TyTree Type -> Constraint
class Decide strat elem typ t | strat t -> typ where
  evidence' :: forall {proxy :: forall k. k -> Type}. proxy strat -> proxy elem -> Path typ t

instance HasField' 'DFS typ ('TyNode typ (t : ts)) where
  evidence _ = Here

instance HasField' 'BFS typ ('TyNode typ (t : ts)) where
  evidence _ = Here

instance HasField' 'DFS typ ('TyNode typ '[]) where
  evidence _ = Here

instance HasField' 'BFS typ ('TyNode typ '[]) where
  evidence _ = Here

instance
  {-# OVERLAPPABLE #-}
  Decide
    'BFS
    (AnyElem typ ts)
    typ
    ('TyNode typ' (t : ts))
  => HasField' 'BFS typ ('TyNode typ' (t : ts))
  where
  evidence _ = evidence' @'BFS @(AnyElem typ ts) @typ @('TyNode typ' (t : ts)) Proxy Proxy

instance
  {-# OVERLAPPABLE #-}
  Decide
    'DFS
    (Not (Elem typ t))
    typ
    ('TyNode typ' (t : ts))
  => HasField' 'DFS typ ('TyNode typ' (t : ts))
  where
  evidence _ = evidence' @'DFS @(Not (Elem typ t)) @typ @('TyNode typ' (t : ts)) Proxy Proxy

instance
  HasField' strat typ t
  => Decide strat 'False typ ('TyNode typ' (t : ts'))
  where
  evidence' _ _ = Deeper (evidence @strat @typ @t Proxy)

instance
  HasField' strat typ ('TyNode typ' ts)
  => Decide strat 'True typ ('TyNode typ' (t' : ts))
  where
  evidence' _ _ = Farther (evidence @strat @typ @('TyNode typ' ts) Proxy)

infixr 4 `HNodeL`

infixr 4 `TyNodeL`