{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Language.Fiddle.Ast.Internal.SyntaxTree
  ( -- Type Families
    NumberType,
    ImportType,
    -- Witness Types
    Witness (..),
    WitnessType,
    -- AST Types
    Name (..),
    Directive (..),
    DirectiveBody (..),
    DirectiveElement (..),
    DirectiveExpression (..),
    Directed (..),
    FiddleUnit (..),
    Identifier (..),
    Expression (..),
    ImportStatement (..),
    ImportList (..),
    FiddleDecl (..),
    ObjTypeBody (..),
    ObjType (..),
    ObjTypeDecl (..),
    Modifier (..),
    ModifierKeyword (..),
    DeferredRegisterBody (..),
    BodyType (..),
    RegisterBody (..),
    RegisterBitsDecl (..),
    RegisterBitsTypeRef (..),
    AnonymousBitsType (..),
    BitType (..),
    EnumBody (..),
    EnumConstantDecl (..),
    PackageBody (..),
    -- Helper Functions
    mapDirected,
    mapDirectedM,
    asDirected,
    undirected,
    -- Utility Functions
    squeeze,
  )
where

import Data.Coerce
import Data.Functor.Identity
import Data.Kind (Type)
import Data.List.NonEmpty
import Data.Proxy
import Data.Text (Text)
import Data.Traversable
import Data.Type.Bool
import Data.Type.Equality
import Data.Typeable
import Data.Void (Void, absurd)
import GHC.Generics
import GHC.TypeLits
import Language.Fiddle.Ast.Internal.Generic
import Language.Fiddle.Ast.Internal.Instances
import Language.Fiddle.Ast.Internal.Kinds
import Language.Fiddle.Ast.Internal.Stage

-- The Type of number during each stage of compilation. When in the first stage,
-- numbers are just strings like anything else. In later stages, numbers get
-- parsed into actual integers. This makes it easier to process later.
type family NumberType (a :: Stage) :: Type where
  NumberType Stage1 = Text
  NumberType Stage2 = Integer
  NumberType s = NumberType (PreviousStage s)

-- The type that represents an import statement. In the early stages of
-- compilation, this is just a string representing the import path, but in later
-- stages of compilation, this actually gets replaced by an abstract
-- representation of the imported material.
type family ImportType (stage :: Stage) :: SynTree where
  ImportType Stage1 = ImportStatement
  ImportType Stage2 = ImportStatement
  ImportType Stage3 = ImportStatement

-- A way to disable or enable a subtree type based on a type-level boolean.
--
-- This is used over GADT's specific parameterization to allow for deriving
-- generics and reduce boilerplate.
--
-- This is a wrapper type to allow defining instances.
data Witness (s :: Bool) where
  Witness :: (WitnessType s) -> Witness s

-- | If the type level expression is true, the witness type is the Unit type, if
-- it is false, the witness type is Void, thus making whatever it touches
-- uninhabitable. This is how parts of the AST is enabled or disabled during
-- compilation.
type family WitnessType (s :: Bool) where
  WitnessType True = ()
  WitnessType False = Void

-- A Name is multiple identifiers separated by dots. It's the way of namespaces
-- to different packages.
data Name :: SynTree where
  Name :: NonEmpty (Identifier f a) -> a -> Name f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents a directive in the Fiddle language. A directive provides 
--   additional metadata or instructions that the compiler can use during 
--   code generation. Directives can be attached to many elements in the 
--   syntax tree.
data Directive :: SynTree where
  Directive ::
    { directiveBody :: f (DirectiveBody f a), -- ^ The body of the directive.
      directiveAnnot :: a -- ^ Annotation for the directive.
    } -> Directive f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents the body of a directive, which consists of multiple elements.
data DirectiveBody :: SynTree where
  DirectiveBody ::
    { directiveElements :: [DirectiveElement f a], -- ^ Elements of the directive.
      directiveBodyAnnot :: a -- ^ Annotation for the directive body.
    } -> DirectiveBody f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents an element in a directive. Can be either a key or a key-value 
--   pair.
data DirectiveElement :: SynTree where
  -- | A simple directive element with a key. The mere presence of this key
  --   holds semantic value.
  DirectiveElementKey ::
    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
      directiveKey :: Identifier f a, -- ^ The key of the directive.
      directiveKeyAnnot :: a -- ^ Annotation for the directive element.
    } -> DirectiveElement f a
  -- | A more complex directive element with a key-value pair, optionally
  --   specifying a backend.
  DirectiveElementKeyValue ::
    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
      directiveKey :: Identifier f a, -- ^ The key of the directive.
      directiveValue :: DirectiveExpression f a, -- ^ The value of the directive.
      directiveKeyValueAnnot :: a -- ^ Annotation for the key-value directive.
    } -> DirectiveElement f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents expressions that can be used within a directive, either a 
--   string or a number.
data DirectiveExpression f a where
  DirectiveString :: 
    { directiveStringValue :: Text, -- ^ String value of the directive.
      directiveStringAnnot :: a -- ^ Annotation for the directive string.
    } -> DirectiveExpression f a
  DirectiveNumber :: 
    { directiveNumberValue :: Text, -- ^ Number value of the directive.
      directiveNumberAnnot :: a -- ^ Annotation for the directive number.
    } -> DirectiveExpression f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | A type that wraps another syntax tree and applies a list of directives to 
--   it.
data Directed t stage f a where
  Directed ::
    { directedDirectives :: [Directive f a], -- ^ List of directives.
      directedSubtree :: t stage f a, -- ^ The wrapped syntax tree.
      directedAnnot :: a -- ^ Annotation for the directed subtree.
    } -> Directed t stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Apply a function to the underlying subtree in a 'Directed' type.
mapDirected :: (t s f a -> t' s' f a) -> Directed t s f a -> Directed t' s' f a
mapDirected fn (Directed dr tfa a) = Directed dr (fn tfa) a

-- | Apply a monadic function to the underlying subtree in a 'Directed' type.
mapDirectedM ::
  (Monad m) => 
  (t s f a -> m (t' s' f a)) -> 
  Directed t s f a -> 
  m (Directed t' s' f a)
mapDirectedM fn (Directed dr tfa a) = Directed dr <$> fn tfa <*> pure a

-- | Convert an annotated syntax tree element to a 'Directed' type with 
--   an empty directive list.
asDirected :: (Annotated (t s)) => t s f a -> Directed t s f a
asDirected tfa = Directed [] tfa (annot tfa)

-- | Extract the underlying subtree from a 'Directed' type, discarding any 
--   directives.
undirected :: Directed t s f a -> t s f a
undirected (Directed _ tfa _) = tfa

-- | The root of the parse tree, containing a list of top-level declarations.
data FiddleUnit (stage :: Stage) (f :: Type -> Type) a where
  FiddleUnit :: 
    { fiddleDecls :: [Directed FiddleDecl stage f a], -- ^ List of declarations.
      fiddleUnitAnnot :: a -- ^ Annotation for the 'FiddleUnit'.
    } -> FiddleUnit stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (FiddleUnit stage)

-- | Represents an identifier with an associated annotation.
data Identifier f a = Identifier 
  { identifierName :: !Text, -- ^ The name of the identifier.
    identifierAnnot :: a -- ^ Annotation for the identifier.
  }
  deriving (Generic, Annotated, Alter, Typeable)

-- | Expressions used within Fiddle, including literals and variables.
data Expression (s :: Stage) :: SynTree where
  -- | A numeric literal, whose value is dependent on the compilation stage.
  LitNum :: 
    { litNumValue :: NumberType stage, -- ^ The numeric value.
      litNumAnnot :: a -- ^ Annotation for the literal.
    } -> Expression stage f a
  -- | A variable reference.
  Var :: 
    { varIdentifier :: Identifier f a, -- ^ The identifier of the variable.
      varAnnot :: a -- ^ Annotation for the variable.
    } -> Expression stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents an import statement in the Fiddle language.
data ImportStatement f a where
  ImportStatement :: 
    { importPath :: Text, -- ^ The path to import.
      importList :: Maybe (ImportList f a), -- ^ Optional list of imported items.
      importStatementAnnot :: a -- ^ Annotation for the import statement.
    } -> ImportStatement f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | A list of imported identifiers.
data ImportList f a where
  ImportList :: 
    { importIdentifiers :: [Identifier f a], -- ^ The list of identifiers.
      importListAnnot :: a -- ^ Annotation for the import list.
    } -> ImportList f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents top-level declarations in Fiddle.
data FiddleDecl :: StagedSynTree where
  -- | An option declaration in the form 'option <ident> <ident>'.
  OptionDecl ::
    { optionKey :: Identifier f a, -- ^ The key of the option.
      optionValue :: Identifier f a, -- ^ The value of the option.
      optionAnnot :: a -- ^ Annotation for the option declaration.
    } -> FiddleDecl stage f a
  -- | An import declaration.
  ImportDecl ::
    { importType :: ImportType stage f a, -- ^ The imported type.
      importDeclAnnot :: a -- ^ Annotation for the import declaration.
    } -> FiddleDecl stage f a
  -- | A using declaration.
  UsingDecl ::
    { usingName :: Name f a, -- ^ The name being used.
      usingAnnot :: a -- ^ Annotation for the using declaration.
    } -> FiddleDecl stage f a
  -- | A package declaration.
  PackageDecl ::
    { packageName :: Name f a, -- ^ The package name.
      packageBody :: f (PackageBody stage f a), -- ^ The body of the package.
      packageAnnot :: a -- ^ Annotation for the package declaration.
    } -> FiddleDecl stage f a
  -- | A location declaration in the form 'location <identifier> = <expr>'.
  LocationDecl ::
    { locationIdent :: Identifier f a, -- ^ The location identifier.
      locationExpr :: Expression stage f a, -- ^ The associated expression.
      locationAnnot :: a -- ^ Annotation for the location declaration.
    } -> FiddleDecl stage f a
  -- | A bits declaration in the form 'bits <identifier> : <type>'.
  BitsDecl ::
    { bitsIdent :: Identifier f a, -- ^ The identifier of the bits.
      bitsType :: BitType stage f a, -- ^ The type of the bits.
      bitsAnnot :: a -- ^ Annotation for the bits declaration.
    } -> FiddleDecl stage f a
  -- | An object type declaration.
  ObjTypeDecl ::
    { objTypeIdent :: Identifier f a, -- ^ The identifier of the object type.
      objTypeBody :: f (ObjTypeBody stage f a), -- ^ The body of the object type.
      objTypeAnnot :: a -- ^ Annotation for the object type declaration.
    } -> FiddleDecl stage f a
  -- | An object declaration in the form 'object <ident> at <expr> : <type>'.
  ObjectDecl ::
    { objectIdent :: Identifier f a, -- ^ The identifier of the object.
      objectLocation :: Expression stage f a, -- ^ The location expression.
      objectType :: ObjType stage f a, -- ^ The type of the object.
      objectAnnot :: a -- ^ Annotation for the object declaration.
    } -> FiddleDecl stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (FiddleDecl stage)

-- | Represents the body of an object type, containing a body type (struct or
--   union), a list of object declarations, and an annotation.
data ObjTypeBody (stage :: Stage) (f :: Type -> Type) a where
  ObjTypeBody ::
    { objBodyType :: BodyType f a, -- ^ The body type (struct or union).
      objBodyDecls :: [Directed ObjTypeDecl stage f a], -- ^ Object declarations.
      objBodyAnnot :: a -- ^ Annotation for the object type body.
    } -> ObjTypeBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents an object type, which can be anonymous, an array, or a
--   reference to another type.
data ObjType stage f a where
  -- | An anonymous object type, allowed only in Stage1. 
  AnonymousObjType ::
    { anonWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
      anonBody :: f (ObjTypeBody stage f a), -- ^ The body of the anonymous type.
      anonAnnot :: a -- ^ Annotation for the anonymous type.
    } -> ObjType stage f a
  -- | An array of object types.
  ArrayObjType ::
    { arrayObjType :: ObjType stage f a, -- ^ The type of the array elements.
      arraySize :: Expression stage f a, -- ^ The size of the array.
      arrayAnnot :: a -- ^ Annotation for the array type.
    } -> ObjType stage f a
  -- | A reference to an existing type by name.
  ReferencedObjType ::
    { refName :: Name f a, -- ^ The name of the referenced type.
      refAnnot :: a -- ^ Annotation for the referenced type.
    } -> ObjType stage f a
  deriving (Typeable, Generic, Alter, Annotated, Typeable)

-- | Represents a declaration inside an object type, such as a register, an
--   assertion, or a substructure.
data ObjTypeDecl stage f a where
  -- | An assertion statement for a specific position.
  AssertPosStatement ::
    { assertWitness :: Witness (StageLessThan stage 3), -- ^ Witness for stage constraint.
      assertExpr :: Expression stage f a, -- ^ The expression for the assertion.
      assertAnnot :: a -- ^ Annotation for the assertion.
    } -> ObjTypeDecl stage f a
  -- | A register declaration.
  RegisterDecl ::
    { regModifier :: Maybe (Modifier f a), -- ^ Optional register modifier.
      regIdent :: Maybe (Identifier f a), -- ^ Optional register identifier.
      regSize :: Expression stage f a, -- ^ Size of the register.
      regBody :: Maybe (RegisterBody stage f a), -- ^ Optional register body.
      regAnnot :: a -- ^ Annotation for the register declaration.
    } -> ObjTypeDecl stage f a
  -- | A reserved declaration for padding or alignment.
  ReservedDecl ::
    { reservedExpr :: Expression stage f a, -- ^ The expression for reserved space.
      reservedAnnot :: a -- ^ Annotation for the reserved declaration.
    } -> ObjTypeDecl stage f a
  -- | A declaration for a substructure (struct or union).
  TypeSubStructure ::
    { subStructureBody :: f (ObjTypeBody stage f a), -- ^ The body of the substructure.
      subStructureName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
      subStructureAnnot :: a -- ^ Annotation for the substructure.
    } -> ObjTypeDecl stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents a modifier for registers (e.g., read-only, read-write).
data Modifier f a where
  ModifierKeyword ::
    { modifierKey :: ModifierKeyword, -- ^ The keyword for the modifier.
      modifierAnnot :: a -- ^ Annotation for the modifier.
    } -> Modifier f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Enumerates the different types of register modifiers.
data ModifierKeyword = Rw | Ro | Wo
  deriving (Eq, Ord, Show, Read, Typeable)

-- | Represents a deferred register body, consisting of a list of bit
--   declarations.
data DeferredRegisterBody stage f a where
  DeferredRegisterBody ::
    { deferredBits :: [Directed RegisterBitsDecl stage f a], -- ^ Bit declarations.
      deferredAnnot :: a -- ^ Annotation for the deferred register body.
    } -> DeferredRegisterBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents the body type (struct or union) in an object.
data BodyType (f :: Type -> Type) a where
  Union ::
    { unionAnnot :: a -- ^ Annotation for the union.
    } -> BodyType f a
  Struct ::
    { structAnnot :: a -- ^ Annotation for the struct.
    } -> BodyType f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents a register body with a body type and deferred bit declarations.
data RegisterBody stage f a where
  RegisterBody ::
    { regBodyType :: BodyType f a, -- ^ The body type of the register.
      regDeferredBody :: f (DeferredRegisterBody stage f a), -- ^ Deferred body.
      regBodyAnnot :: a -- ^ Annotation for the register body.
    } -> RegisterBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents declarations within a register, such as defined bits,
--   reserved bits, or substructures.
data RegisterBitsDecl stage f a where
  -- | Declaration for reserved bits.
  ReservedBits ::
    { reservedBitsExpr :: Expression stage f a, -- ^ Expression for reserved bits.
      reservedBitsAnnot :: a -- ^ Annotation for the reserved bits.
    } -> RegisterBitsDecl stage f a
  -- | Declaration for defined bits in a register.
  DefinedBits ::
    { definedBitsModifier :: Maybe (Modifier f a), -- ^ Optional modifier for the bits.
      definedBitsIdent :: Identifier f a, -- ^ Identifier for the bits.
      definedBitsTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Type reference for the bits.
      definedBitsAnnot :: a -- ^ Annotation for the defined bits.
    } -> RegisterBitsDecl stage f a
  -- | Substructure within a register.
  BitsSubStructure ::
    { bitsSubRegBody :: RegisterBody stage f a, -- ^ The body of the substructure.
      bitsSubName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
      bitsSubAnnot :: a -- ^ Annotation for the substructure.
    } -> RegisterBitsDecl stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents different ways to refer to register bits, either as an array,
--   a reference to a type, an anonymous type, or just bits.
data RegisterBitsTypeRef stage f a where
  -- | An array of bits with a specified size.
  RegisterBitsArray ::
    { bitsArrayTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Reference to the array type.
      bitsArraySize :: Expression stage f a, -- ^ Size of the array.
      bitsArrayAnnot :: a -- ^ Annotation for the array.
    } -> RegisterBitsTypeRef stage f a
  -- | A reference to another type by name.
  RegisterBitsReference ::
    { bitsRefName :: Name f a, -- ^ The name of the referenced type.
      bitsRefAnnot :: a -- ^ Annotation for the reference.
    } -> RegisterBitsTypeRef stage f a
  -- | An anonymous type for register bits, used in Stage1.
  RegisterBitsAnonymousType ::
    { anonBitsWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
      anonBitsType :: AnonymousBitsType stage f a, -- ^ The anonymous type.
      anonBitsAnnot :: a -- ^ Annotation for the anonymous type.
    } -> RegisterBitsTypeRef stage f a
  -- | A direct specification of bits as an expression.
  RegisterBitsJustBits ::
    { justBitsExpr :: Expression stage f a, -- ^ Expression for the bits.
      justBitsAnnot :: a -- ^ Annotation for the bits.
    } -> RegisterBitsTypeRef stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents an anonymous bit type, such as an enum, used in Stage1.
data AnonymousBitsType stage f a where
  AnonymousEnumBody ::
    { anonEnumExpr :: Expression stage f a, -- ^ Expression defining the enum size.
      anonEnumBody :: f (EnumBody stage f a), -- ^ The body of the enum.
      anonEnumAnnot :: a -- ^ Annotation for the anonymous enum.
    } -> AnonymousBitsType stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents a bit type, either an enumeration or raw bits.
data BitType (stage :: Stage) (f :: Type -> Type) a where
  -- | An enumeration type for bits.
  EnumBitType ::
    { enumBitExpr :: Expression stage f a, -- ^ Expression defining the enum size.
      enumBitBody :: f (EnumBody stage f a), -- ^ The body of the enum.
      enumBitAnnot :: a -- ^ Annotation for the enumeration.
    } -> BitType stage f a
  -- | A raw bit type.
  RawBits ::
    { rawBitsExpr :: Expression stage f a, -- ^ Expression defining the bits.
      rawBitsAnnot :: a -- ^ Annotation for the raw bits.
    } -> BitType stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents the body of an enumeration.
data EnumBody (stage :: Stage) (f :: Type -> Type) a where
  EnumBody ::
    { enumConsts :: [Directed EnumConstantDecl stage f a], -- ^ Enum constant declarations.
      enumBodyAnnot :: a -- ^ Annotation for the enum body.
    } -> EnumBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents a declaration for an enumeration constant.
data EnumConstantDecl stage f a where
  -- | A named constant in the enum.
  EnumConstantDecl ::
    { enumConstIdent :: Identifier f a, -- ^ Identifier for the constant.
      enumConstExpr :: Expression stage f a, -- ^ Expression defining the constant.
      enumConstAnnot :: a -- ^ Annotation for the constant.
    } -> EnumConstantDecl stage f a
  -- | A reserved value in the enum.
  EnumConstantReserved ::
    { enumReservedExpr :: Expression stage f a, -- ^ Expression for the reserved value.
      enumReservedAnnot :: a -- ^ Annotation for the reserved value.
    } -> EnumConstantDecl stage f a
  deriving (Generic, Annotated, Alter, Typeable)

-- | Represents the body of a package, containing a list of declarations.
data PackageBody (stage :: Stage) (f :: Type -> Type) a where
  PackageBody ::
    { packageBodyDecls :: [Directed FiddleDecl stage f a], -- ^ Declarations in the package.
      packageBodyAnnot :: a -- ^ Annotation for the package body.
    } -> PackageBody stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (PackageBody stage)


squeeze :: (Alter t, Traversable f, Monad f) => t f a -> f (t Identity a)
squeeze = alter (fmap Identity) return