module Language.Fiddle.Ast.Internal.Instances
  ( module X,
    Alter (..),
    AdvanceStage (..),
    CompilationStage (..),
    Annotated (..),
    GAnnot (..),
    TreeType,
  )
where

import Data.Functor.Identity
import Data.Kind
import Data.Typeable
import GHC.Generics
import Language.Fiddle.Ast.Internal.Instances.Walk as X
import Language.Fiddle.Ast.Internal.Kinds
import Language.Fiddle.Ast.Internal.Stage

-- Class for walking a syntax tree under the context of a monad and modifying
-- the different parts of the SynTree type..
class Alter (t :: SynTree) where
  alter ::
    forall f1 f2 a1 a2 m.
    (Monad m, Traversable f1) =>
    (forall z. f1 z -> m (f2 z)) ->
    (a1 -> m a2) ->
    t f1 a1 ->
    m (t f2 a2)
  default alter ::
    forall f1 f2 a1 a2 m.
    ( Generic (t f1 a1),
      Generic (t f2 a2),
      Traversable f1,
      GAlter t f1 f2 a1 a2 (Rep (t f1 a1)) (Rep (t f2 a2)),
      Monad m
    ) =>
    ( forall z.
      f1 z ->
      m (f2 z)
    ) ->
    (a1 -> m a2) ->
    t f1 a1 ->
    m (t f2 a2)
  alter ffn fn t = to <$> galter (proxyOf t) ffn fn (from t)

-- | 'CompilationStage' is a type class representing a stage in the compiler
--   pipeline. Each 'stage' has associated types that define how it transforms
--   syntax trees, manages state, and handles annotations.
--
--   This class requires:
--   * A 'StageMonad' for handling monadic actions during the compilation process.
--   * A 'StageFunctor' for traversing the syntax tree.
--   * A 'StageState' to represent local state information during traversal.
class
  (Monad (StageMonad stage), Traversable (StageFunctor stage)) =>
  CompilationStage stage
  where
  -- | The next stage in the compilation pipeline.
  type StageAfter stage :: Stage

  -- | The monadic context for this stage (e.g., 'Compile' or 'Either').
  type StageMonad stage :: Type -> Type

  -- | The state type associated with this stage. This state is designed to be
  --   used as a 'local state', meaning any modifications to this state are
  --   only visible to child nodes during the traversal. To share state across
  --   sibling nodes, use the 'StageMonad'.
  type StageState stage :: Type

  -- | The functor used for traversing the syntax tree during this stage.
  type StageFunctor stage :: Type -> Type

  -- | The type of annotations associated with nodes in the syntax tree.
  type StageAnnotation stage :: Type

-- | Utility type that captures the structure of a syntax tree node for a given
--   stage. This type alias simplifies references to the full tree structure in
--   other parts of the code.
type TreeType (t :: StagedSynTree) (s :: Stage) =
  t s (StageFunctor s) (StageAnnotation s)

class
  (CompilationStage stage) =>
  StageConvertible stage from to
  where
  convertInStage :: proxy stage -> StageState stage -> from -> StageMonad stage to

instance (CompilationStage s, Applicative (StageMonad s)) => StageConvertible s a a where
  convertInStage _ _ = pure

-- | 'AdvanceStage' defines how to transform an Abstract Syntax Tree (AST) node
--   from one stage to the next in the compiler pipeline. This transformation
--   can be customized per node type, or a default generic implementation can be
--   used.
class
  (CompilationStage stage) =>
  AdvanceStage (stage :: Stage) (t :: StagedSynTree)
  where
  -- | 'advanceStage' transitions a syntax tree node from the current stage to
  --   the next stage ('StageAfter stage'). It takes the current local state
  --   ('StageState') and the syntax tree node ('TreeType') as input and returns
  --   a monadic action that produces the transformed node for the next stage.
  --
  --   This function typically performs tree transformations, checks, or other
  --   modifications necessary for the compilation process.
  --
  --   Parameters:
  --   * 'StageState stage' - The local state for this stage of the compiler.
  --     This state is only visible to child nodes. Any changes do not affect
  --     sibling nodes.
  --   * 'TreeType t stage' - The syntax tree node at the current stage.
  --
  --   Returns: A monadic action that produces the transformed syntax tree node
  --   at the next stage.
  advanceStage ::
    StageState stage -> -- Local state for the current stage
    TreeType t stage -> -- Syntax tree node at the current stage
    StageMonad
      stage
      ( t
          (StageAfter stage) -- The next stage in the pipeline
          (StageFunctor stage) -- The functor associated with this stage
          (StageAnnotation stage) -- Annotation type for this stage
      )

  -- | Default implementation of 'advanceStage' using generics. This leverages
  --   'GAdvanceStage' to automatically traverse and transform the syntax tree.
  --   If an AST type derives 'Generic', this default can be used to reduce
  --   boilerplate code.
  --
  --   Before performing the generic transformation with 'gAdvanceStage', the
  --   'modifyState' function is called to potentially alter the local state
  --   based on the current node.
  --
  --   This implementation is useful for cases where the tree structure remains
  --   mostly unchanged between stages, and only the state or annotations need
  --   to be adjusted.
  default advanceStage ::
    ( GAdvanceStage
        stage
        (StageState stage) -- The local state for this stage
        (StageMonad stage) -- The monadic context of this stage
        (Rep (TreeType t stage)) -- Generic representation of the current tree type
        (Rep (t (StageAfter stage) (StageFunctor stage) (StageAnnotation stage))), -- Generic representation of the next stage's tree type
      Generic (TreeType t stage), -- The current tree type must be an instance of 'Generic'
      Generic
        ( t
            (StageAfter stage) -- The tree type at the next stage
            (StageFunctor stage) -- The functor for the next stage
            (StageAnnotation stage) -- The annotation type for the next stage
        )
    ) =>
    StageState stage -> -- Local state for the current stage
    TreeType t stage -> -- Syntax tree node at the current stage
    StageMonad
      stage
      ( t
          (StageAfter stage) -- The tree type for the next stage
          (StageFunctor stage) -- The functor for the next stage
          (StageAnnotation stage) -- The annotation type for the next stage
      )
  advanceStage s t = do
    -- Modify the local state for this node before performing the transformation
    s' <- modifyState t s
    specific <- customAdvanceStage t s

    case specific of
      Nothing ->
        -- Perform the generic transformation using 'gAdvanceStage'
        to <$> gAdvanceStage (Proxy :: Proxy stage) s' (from t)
      Just ast -> return ast

  -- | 'modifyState' allows for changes to the local state ('StageState') before
  --   transforming the syntax tree node. This is called on each node during the
  --   traversal, allowing the state to be adjusted based on the current node's
  --   information.
  --
  --   The default implementation simply returns the unchanged state, but it can
  --   be overridden to implement custom state modifications.
  --
  --   Note: This state modification is local to the current node and its
  --   children. Changes to this state are not visible to sibling nodes. If
  --   information needs to be shared across siblings, consider using the
  --   'StageMonad' for that purpose.
  --
  --   'modifyState' is only called by the default implementation of
  --   advanceStage. If 'advaceStage' is overridden, then overriding this
  --   function is superfluous.
  --
  --   Parameters:
  --   * 'TreeType t stage' - The syntax tree node at the current stage.
  --   * 'StageState stage' - The local state for this stage.
  --
  --   Returns: A monadic action that produces the (potentially modified) state.
  modifyState ::
    TreeType t stage -> -- Syntax tree node at the current stage
    StageState stage -> -- Local state for the current stage
    StageMonad stage (StageState stage) -- The modified local state
  modifyState _ = return

  -- \| An optional method that allows for partial customization of the
  -- 'advanceStage' process. This method can be used to handle specific cases
  -- in the input without requiring the implementer to redefine 'advanceStage'
  -- for the entire structure. If this method returns 'Nothing', the default
  -- generic implementation of 'advanceStage' is used.
  --
  -- This method is useful when only certain conditions or patterns in the
  -- syntax tree need special handling during the stage transition. Implementers
  -- can focus on those specific conditions and leave the remaining cases to the
  -- generic traversal.
  --
  -- Parameters:
  --   - 'StageState stage': The local state for the current stage.
  --   - 'TreeType t stage': The syntax tree node at the current stage.
  --
  -- Returns:
  --   - 'StageMonad stage (Maybe (t (StageAfter stage) (StageFunctor stage)
  --     (StageAnnotation stage)))': A monadic computation that either returns
  --     'Just' the transformed tree node for the next stage, or 'Nothing' to
  --     continue with the default generic implementation.
  customAdvanceStage ::
    TreeType t stage -> -- Syntax tree node at the current stage
    StageState stage -> -- Local state for the current stage
    StageMonad
      stage
      ( Maybe
          ( t
              (StageAfter stage) -- The next stage in the pipeline
              (StageFunctor stage) -- Functor associated with the next stage
              (StageAnnotation stage) -- Annotation type for the next stage
          )
      )
  customAdvanceStage _ _ = return Nothing

-- | 'GAdvanceStage' is a helper type class that performs the transformation
--   of the generic representation of a syntax tree node. It is used by the
--   default implementation of 'advanceStage' to traverse and modify nodes
--   automatically.
class GAdvanceStage (stage :: Stage) s m from to where
  gAdvanceStage :: Proxy stage -> s -> from x -> m (to x)

-- A syntax tree object is annotated if it has an annotation 'a' as the last
-- element.
class Annotated (t :: SynTree) where
  annot :: t f a -> a
  setAnnot :: (a -> a) -> t f a -> t f a

  default annot :: (Generic (t f a), GAnnot a (Rep (t f a))) => t f a -> a
  annot t = gannot (from t)

  default setAnnot :: (Generic (t f a), GAnnot a (Rep (t f a))) => (a -> a) -> t f a -> t f a
  setAnnot f t = to $ gsetAnnot f (from t)

-- Generic implementations of common typeclass for SyntaxTrees.
--
-- This is where we try to hide the pig behind the curtain.
--
--
-- ---------------------------------------------------------------
--
--
--
class GAlter t f1 f2 a1 a2 r1 r2 where
  galter ::
    forall proxy x m.
    (Monad m, Traversable f1) =>
    proxy t ->
    (forall z. f1 z -> m (f2 z)) ->
    (a1 -> m a2) ->
    r1 x ->
    m (r2 x)

{- Altering a record with type a1 will apply the mapping function and produce a
    record with type a2 -}
instance GAlter t f1 f2 a1 a2 (Rec0 a1) (Rec0 a2) where
  galter _ _ fn k1 = K1 <$> fn (unK1 k1)

{- Base-Case. Altering unrelated leaf types will do nothing. -}
instance GAlter t f1 f2 a1 a2 (Rec0 u1) (Rec0 u1) where
  galter _ _ _ = return

{- Recursive case. Call alter on sub-structures. -}
instance
  (Alter u) =>
  GAlter t f1 f2 a1 a2 (Rec0 (u f1 a1)) (Rec0 (u f2 a2))
  where
  galter _ ffn fn k1 =
    K1 <$> alter ffn fn (unK1 k1)

{- Recursive case. Called when there are list of substructures that need to be
    recused. -}
instance
  (Alter u, Traversable l) =>
  GAlter t f1 f2 a1 a2 (Rec0 (l (u f1 a1))) (Rec0 (l (u f2 a2)))
  where
  galter _ ffn fn k1 =
    K1 <$> mapM (alter ffn fn) (unK1 k1)

-- instance GAlter t f1 f2 a1 a2 (Rec0 (f1 z)) (Rec0 (f2 z)) where
--   galter _ ffn _ k1 = K1 <$> ffn (unK1 k1)

{- Generic altering. Descends into the function and alters whatever is inside
    the functor and then transforms the functor using the ffn function. -}
instance
  (Traversable f1, Alter u) =>
  GAlter t f1 f2 a1 a2 (Rec0 (f1 (u f1 a1))) (Rec0 (f2 (u f2 a2)))
  where
  galter _ ffn fn k1 = do
    newK <- mapM (alter ffn fn) (unK1 k1)
    K1 <$> ffn newK

instance
  (Traversable f1) =>
  GAlter t f1 f2 a1 a2 (Rec0 (f1 x)) (Rec0 (f2 x))
  where
  galter _ ffn _ (K1 k) = K1 <$> ffn k

instance
  ( GAlter t f1 f2 a1 a2 l1 l2,
    GAlter t f1 f2 a1 a2 r1 r2
  ) =>
  GAlter t f1 f2 a1 a2 (l1 :*: r1) (l2 :*: r2)
  where
  galter proxy ffn fn (a :*: b) = do
    a' <- galter proxy ffn fn a
    b' <- galter proxy ffn fn b
    return (a' :*: b')

instance
  ( GAlter t f1 f2 a1 a2 l1 l2,
    GAlter t f1 f2 a1 a2 r1 r2
  ) =>
  GAlter t f1 f2 a1 a2 (l1 :+: r1) (l2 :+: r2)
  where
  galter proxy ffn fn (R1 r) = R1 <$> galter proxy ffn fn r
  galter proxy ffn fn (L1 l) = L1 <$> galter proxy ffn fn l

instance
  (GAlter t f1 f2 a1 a2 r1 r2) =>
  GAlter t f1 f2 a1 a2 (M1 i c r1) (M1 i c r2)
  where
  galter proxy ffn fn (M1 a) = M1 <$> galter proxy ffn fn a

class GAnnot a r where
  gannot :: r x -> a
  gsetAnnot :: (a -> a) -> r x -> r x

instance GAnnot a (Rec0 a) where
  gannot = unK1
  gsetAnnot fn (K1 t) = K1 (fn t)

instance (GAnnot a r) => GAnnot a (l :*: r) where
  gannot (_ :*: r) = gannot r
  gsetAnnot fn (l :*: r) = l :*: gsetAnnot fn r

instance (GAnnot a r, GAnnot a l) => GAnnot a (l :+: r) where
  gannot (R1 r) = gannot r
  gannot (L1 l) = gannot l
  gsetAnnot fn (R1 r) = R1 (gsetAnnot fn r)
  gsetAnnot fn (L1 l) = L1 (gsetAnnot fn l)

instance (GAnnot a r) => GAnnot a (M1 i c r) where
  gannot (M1 a) = gannot a
  gsetAnnot fn (M1 a) = M1 (gsetAnnot fn a)

proxyOf :: t f a -> Proxy t
proxyOf _ = Proxy

instance (Alter t, Traversable f) => Functor (t f) where
  fmap f t = runIdentity (alter return (return . f) t)

-- | 'GAdvanceStage' instance for metadata wrappers ('M1').
--   This instance allows advancing the stage of a metadata node in a generic
--   representation. The metadata node ('M1') wraps another node ('s1'), which
--   is recursively advanced to the next stage using 'gAdvanceStage'.
instance
  (Monad m, GAdvanceStage stage s m s1 s2) =>
  GAdvanceStage stage s m (M1 i c s1) (M1 i c s2)
  where
  gAdvanceStage pxy s (M1 a) = M1 <$> gAdvanceStage pxy s a

-- | 'GAdvanceStage' instance for sum types (':+:'). This handles the case
--   where the generic representation of a type is a sum (i.e., an 'Either'-like
--   choice between two alternatives). The sum type can be either 'L1' (left)
--   or 'R1' (right), and 'gAdvanceStage' is called recursively on the selected
--   branch.
instance
  (Monad m, GAdvanceStage stage s m l1 l2, GAdvanceStage stage s m r1 r2) =>
  GAdvanceStage stage s m (l1 :+: r1) (l2 :+: r2)
  where
  gAdvanceStage pxy s (R1 r) = R1 <$> gAdvanceStage pxy s r
  gAdvanceStage pxy s (L1 l) = L1 <$> gAdvanceStage pxy s l

-- | 'GAdvanceStage' instance for product types (':*:'). This handles cases
--   where the generic representation of a type is a product (i.e., a tuple of
--   multiple components). It recursively advances each component ('l' and 'r')
--   to the next stage.
instance
  (Monad m, GAdvanceStage stage s m l1 l2, GAdvanceStage stage s m r1 r2) =>
  GAdvanceStage stage s m (l1 :*: r1) (l2 :*: r2)
  where
  gAdvanceStage pxy s (l :*: r) =
    (:*:) <$> gAdvanceStage pxy s l <*> gAdvanceStage pxy s r

-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a single
--   AST element ('t') to be advanced. This instance covers the case where the
--   field is an individual AST node that implements 'AdvanceStage'.
--   It advances this node using 'advanceStage'.
instance
  ( Monad m,
    AdvanceStage stage t',
    Traversable f,
    StageAfter stage ~ stage',
    StageMonad stage ~ m,
    StageState stage ~ s,
    StageFunctor stage ~ f,
    StageAnnotation stage ~ a
  ) =>
  GAdvanceStage stage s m (Rec0 (t' stage f a)) (Rec0 (t' stage' f a))
  where
  gAdvanceStage _ st (K1 val) = K1 <$> advanceStage st val

-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a functor
--   ('func') of AST elements ('t'). This handles cases where the field is a
--   container (e.g., list, 'Maybe') of AST nodes that need to be advanced.
--   Each node in the container is transformed using 'advanceStage'.
instance
  ( Monad m,
    AdvanceStage stage t',
    Traversable f,
    StageAfter stage ~ stage',
    StageMonad stage ~ m,
    StageState stage ~ s,
    StageFunctor stage ~ f,
    StageAnnotation stage ~ a,
    Traversable func
  ) =>
  GAdvanceStage stage s m (Rec0 (func (t' stage f a))) (Rec0 (func (t' stage' f a)))
  where
  gAdvanceStage _ st (K1 val) = K1 <$> mapM (advanceStage st) val

-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a
--   functor ('f') wrapping an AST element. This handles cases where the field
--   is a container ('f') of AST nodes that need to be advanced. Each node in
--   the container is transformed using 'advanceStage'.
instance
  ( Monad m,
    AdvanceStage stage t',
    Traversable f,
    StageAfter stage ~ stage',
    StageMonad stage ~ m,
    StageState stage ~ s,
    StageFunctor stage ~ f,
    StageAnnotation stage ~ a
  ) =>
  GAdvanceStage stage s m (Rec0 (f (t' stage f a))) (Rec0 (f (t' stage' f a)))
  where
  gAdvanceStage _ st (K1 val) = K1 <$> mapM (advanceStage st) val

-- | 'GAdvanceStage' instance for simple record fields ('Rec0') that do not
--   need to change between stages. This is used for fields that are not AST
--   nodes and remain the same when advancing the stage (e.g., primitive
--   types like 'Int', 'Bool', etc.).
instance (Monad m) => GAdvanceStage stage s m (Rec0 a) (Rec0 a) where
  gAdvanceStage _ _ (K1 val) = return (K1 val)

-- | 'GAdvanceStage' instance for records which can be converted to eathother
-- for the current stage..
instance
  ( Monad m,
    StageConvertible stage a b,
    StageState stage ~ s,
    StageMonad stage ~ m
  ) =>
  GAdvanceStage stage s m (Rec0 a) (Rec0 b)
  where
  gAdvanceStage pxy s (K1 val) = K1 <$> convertInStage pxy s val