{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ConstrainedClassMethods #-}
module HSVIS.AST where

import Data.Bifunctor (bimap, first, second)
import qualified Data.Kind as DK
import Data.List (intersperse)
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy

import HSVIS.Pretty


-- | An AST type has a single stage argument.
type ASTType = DK.Type -> DK.Type

-- | Trees that Grow extension type family
type family X (f :: ASTType) s

-- | Additional constructors
data family E (f :: ASTType) s

class HasExt f where
  type HasXField f :: Bool
  type HasECon f :: Bool
  type ContainsData f :: [ASTType]
  extOf :: HasXField f ~ 'True => f s -> X f s
  extMap :: Proxy f -> Proxy s1 -> Proxy s2
         -> XMapperTelescope (f ': ContainsData f) s1 s2
              (EMapperTelescope (f ': ContainsData f) s1 s2
                (f s1 -> f s2))

type family If c a b where
  If 'True a b = a
  If 'False a b = b

type family XMapperTelescope fs s1 s2 a where
  XMapperTelescope '[] s1 s2 a = a
  XMapperTelescope (f ': fs) s1 s2 a =
    If (HasXField f)
       ((X f s1 -> X f s2) -> XMapperTelescope fs s1 s2 a)
       (XMapperTelescope fs s1 s2 a)

type family EMapperTelescope fs s1 s2 a where
  EMapperTelescope '[] s1 s2 a = a
  EMapperTelescope (f ': fs) s1 s2 a =
    If (HasECon f)
       (If (HasXField f) (X f s1 -> E f s1 -> f s2)
                         (E f s1 -> f s2)
        -> EMapperTelescope fs s1 s2 a)
       (EMapperTelescope fs s1 s2 a)


newtype Name = Name String
  deriving (Show, Eq, Ord)

data Program s = Program [DataDef s] [FunDef s]
deriving instance (Show (DataDef s), Show (FunDef s)) => Show (Program s)

data DataDef s = DataDef (X DataDef s) Name [(X Type s, Name)] [(Name, [Type s])]
deriving instance (Show (X DataDef s), Show (X Type s), Show (E Type s)) => Show (DataDef s)

data FunDef s = FunDef (X FunDef s) Name (TypeSig s) (NonEmpty (FunEq s))
deriving instance (Show (X FunDef s), Show (TypeSig s), Show (FunEq s)) => Show (FunDef s)

data TypeSig s
  = TypeSig (Type s)
  | TypeSigExt !(E TypeSig s)
deriving instance (Show (E TypeSig s), Show (Type s)) => Show (TypeSig s)

data FunEq s = FunEq (X FunEq s) Name [Pattern s] (RHS s)
deriving instance (Show (X FunEq s), Show (Pattern s), Show (RHS s)) => Show (FunEq s)

data Kind s
  = KType (X Kind s)
  | KFun (X Kind s) (Kind s) (Kind s)

  -- extension point
  | KExt (X Kind s) !(E Kind s)
deriving instance (Show (X Kind s), Show (E Kind s)) => Show (Kind s)
deriving instance (Eq (X Kind s), Eq (E Kind s)) => Eq (Kind s)
deriving instance (Ord (X Kind s), Ord (E Kind s)) => Ord (Kind s)

data Type s
  = TApp (X Type s) (Type s) [Type s]
  | TTup (X Type s) [Type s]
  | TList (X Type s) (Type s)
  | TFun (X Type s) (Type s) (Type s)
  | TCon (X Type s) Name
  | TVar (X Type s) Name
  | TForall (X Type s) Name (Type s)  -- ^ implicit

  -- extension point
  | TExt (X Type s) !(E Type s)
deriving instance (Show (X Type s), Show (E Type s)) => Show (Type s)
deriving instance (Eq (X Type s), Eq (E Type s)) => Eq (Type s)
deriving instance (Ord (X Type s), Ord (E Type s)) => Ord (Type s)

data Pattern s
  = PWildcard (X Pattern s)
  | PVar      (X Pattern s) Name
  | PAs       (X Pattern s) Name (Pattern s)
  | PCon      (X Pattern s) Name [Pattern s]
  | POp       (X Pattern s) (Pattern s) Operator (Pattern s)
  | PList     (X Pattern s) [Pattern s]
  | PTup      (X Pattern s) [Pattern s]
deriving instance (Show (X Pattern s)) => Show (Pattern s)

data RHS s
  = Guarded (X RHS s) [(Expr s, Expr s)]  -- currently not parsed
  | Plain   (X RHS s) (Expr s)
deriving instance (Show (X RHS s), Show (Expr s)) => Show (RHS s)

data Expr s
  = ELit  (X Expr s) Literal
  | EVar  (X Expr s) Name
  | ECon  (X Expr s) Name
  | EList (X Expr s) [Expr s]
  | ETup  (X Expr s) [Expr s]
  | EApp  (X Expr s) (Expr s) [Expr s]
  | EOp   (X Expr s) (Expr s) Operator (Expr s)
  | EIf   (X Expr s) (Expr s) (Expr s) (Expr s)
  | ECase (X Expr s) (Expr s) [(Pattern s, RHS s)]
  -- TODO: pattern bindings?
  | ELet  (X Expr s) [FunDef s] (Expr s)
  | EError (X Expr s)
deriving instance (Show (X Expr s), Show (Pattern s), Show (RHS s), Show (FunDef s)) => Show (Expr s)

data Literal = LInt Integer | LFloat Rational | LChar Char | LString String
  deriving (Show)

data Operator = OAdd | OSub | OMul | ODiv | OMod | OEqu | OPow
              | OCons
  deriving (Show)

instance Pretty Name where
  prettysPrec _ (Name n) = showString n

instance Pretty (E Kind s) => Pretty (Kind s) where
  prettysPrec _ (KType _) = showString "Type"
  prettysPrec d (KFun _ a b) = showParen (d > -1) $
    prettysPrec 0 a . showString " -> " . prettysPrec (-1) b
  prettysPrec d (KExt _ e) = prettysPrec d e

instance Pretty (E Type s) => Pretty (Type s) where
  prettysPrec d (TApp _ a ts) = showParen (d > 10) $
    prettysPrec 10 a . foldr (.) id [showString " " . prettysPrec 11 t | t <- ts]
  prettysPrec _ (TTup _ ts) =
    showString "(" . foldr (.) id (intersperse (showString ",") (map (prettysPrec 0) ts)) . showString ")"
  prettysPrec _ (TList _ t) =
    showString "[" . prettysPrec 0 t . showString "]"
  prettysPrec d (TFun _ a b) = showParen (d > -1) $
    prettysPrec 0 a . showString " -> " . prettysPrec (-1) b
  prettysPrec _ (TCon _ n) = prettysPrec 11 n
  prettysPrec _ (TVar _ n) = prettysPrec 11 n
  prettysPrec d (TForall _ n t) = showParen (d > -1) $
    showString "forall " . prettysPrec 11 n . showString "." . prettysPrec (-1) t
  prettysPrec d (TExt _ e) = prettysPrec d e

instance (Pretty (X Type s), Pretty (E Type s)) => Pretty (DataDef s) where
  prettysPrec _ (DataDef _ name vars cons) =
    showString "data " . prettysPrec 11 name
    . foldr (.) id [showString " (" . prettysPrec 11 n . showString " :: " . prettysPrec 11 k . showString ")"
                   | (k, n) <- vars]
    . showString " = "
    . foldr (.) id (intersperse (showString " | ")
                      [prettysPrec 11 cname . foldr (.) id [showString " " . prettysPrec 11 t | t <- fields]
                      | (cname, fields) <- cons])

instance Pretty Operator where
  prettysPrec _ = showString . \case
    OAdd -> "(+)"
    OSub -> "(-)"
    OMul -> "(*)"
    ODiv -> "(/)"
    OMod -> "(%)"
    OEqu -> "(==)"
    OPow -> "(^)"
    OCons -> "(:)"

instance HasExt DataDef where
  type HasXField DataDef = 'True
  type HasECon DataDef = 'False
  type ContainsData DataDef = '[Type]
  extOf (DataDef x _ _ _) = x
  extMap _ ps1 ps2 fdd fty gty (DataDef x n ps cs) =
    DataDef (fdd x) n (map (first fty) ps) (map (second (map (extMap (Proxy @Type) ps1 ps2 fty gty))) cs)

instance HasExt FunDef where
  type HasXField FunDef = 'True
  type HasECon FunDef = 'False
  type ContainsData FunDef = '[TypeSig, FunEq, Type, Pattern, RHS, Expr]
  extOf (FunDef x _ _ _) = x
  extMap _ ps1 ps2 ffd ffe fty fpa frh fex gts gty (FunDef x n sig es) =
    FunDef (ffd x) n
           (extMap (Proxy @TypeSig) ps1 ps2 fty gts gty sig)
           (fmap (extMap (Proxy @FunEq) ps1 ps2 ffe ffd fty fpa frh fex gts gty) es)

instance HasExt TypeSig where
  type HasXField TypeSig = 'False
  type HasECon TypeSig = 'True
  type ContainsData TypeSig = '[Type]
  extOf _ = undefined
  extMap _ ps1 ps2 fty _   gty (TypeSig ty) =
    TypeSig (extMap (Proxy @Type) ps1 ps2 fty gty ty)
  extMap _ _   _   _   gts _   (TypeSigExt e) = gts e

instance HasExt FunEq where
  type HasXField FunEq = 'True
  type HasECon FunEq = 'False
  type ContainsData FunEq = '[FunDef, TypeSig, Type, Pattern, RHS, Expr]
  extOf (FunEq x _ _ _) = x
  extMap _ ps1 ps2 ffe ffd fty fpa frh fex gts gty (FunEq x n ps rhs) =
    FunEq (ffe x) n (map (extMap (Proxy @Pattern) ps1 ps2 fpa) ps)
          (extMap (Proxy @RHS) ps1 ps2 frh ffd ffe fty fpa fex gts gty rhs)

instance HasExt Kind where
  type HasXField Kind = 'True
  type HasECon Kind = 'True
  type ContainsData Kind = '[]
  extOf (KType x) = x
  extOf (KFun x _ _) = x
  extOf (KExt x _) = x
  extMap _ _   _   f _ (KType x) = KType (f x)
  extMap p ps1 ps2 f g (KFun x a b) = KFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
  extMap _ _   _   _ g (KExt x a) = g x a

instance HasExt Type where
  type HasXField Type = 'True
  type HasECon Type = 'True
  type ContainsData Type = '[]
  extOf (TApp x _ _) = x
  extOf (TTup x _) = x
  extOf (TList x _) = x
  extOf (TFun x _ _) = x
  extOf (TCon x _) = x
  extOf (TVar x _) = x
  extOf (TForall x _ _) = x
  extOf (TExt x _) = x

  extMap p ps1 ps2 f g (TApp x a b) = TApp (f x) (extMap p ps1 ps2 f g a) (map (extMap p ps1 ps2 f g) b)
  extMap p ps1 ps2 f g (TTup x as) = TTup (f x) (map (extMap p ps1 ps2 f g) as)
  extMap p ps1 ps2 f g (TList x a) = TList (f x) (extMap p ps1 ps2 f g a)
  extMap p ps1 ps2 f g (TFun x a b) = TFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
  extMap _ _   _   f _ (TCon x n) = TCon (f x) n
  extMap _ _   _   f _ (TVar x n) = TVar (f x) n
  extMap p ps1 ps2 f g (TForall x n a) = TForall (f x) n (extMap p ps1 ps2 f g a)
  extMap _ _   _   _ g (TExt x e) = g x e

instance HasExt Pattern where
  type HasXField Pattern = 'True
  type HasECon Pattern = 'False
  type ContainsData Pattern = '[]
  extOf (PWildcard x) = x
  extOf (PVar x _) = x
  extOf (PAs x _ _) = x
  extOf (PCon x _ _) = x
  extOf (POp x _ _ _) = x
  extOf (PList x _) = x
  extOf (PTup x _) = x

  extMap _ _   _   f (PWildcard x) = PWildcard (f x)
  extMap _ _   _   f (PVar x n) = PVar (f x) n
  extMap p ps1 ps2 f (PAs x n a) = PAs (f x) n (extMap p ps1 ps2 f a)
  extMap p ps1 ps2 f (PCon x n as) = PCon (f x) n (map (extMap p ps1 ps2 f) as)
  extMap p ps1 ps2 f (POp x a n b) = POp (f x) (extMap p ps1 ps2 f a) n (extMap p ps1 ps2 f b)
  extMap p ps1 ps2 f (PList x as) = PList (f x) (map (extMap p ps1 ps2 f) as)
  extMap p ps1 ps2 f (PTup x as) = PTup (f x) (map (extMap p ps1 ps2 f) as)

instance HasExt RHS where
  type HasXField RHS = 'True
  type HasECon RHS = 'False
  type ContainsData RHS = '[FunDef, TypeSig, FunEq, Type, Pattern, Expr]
  extOf (Guarded x _) = x
  extOf (Plain x _) = x

  extMap _ ps1 ps2 frh ffd ffe fty fpa fex gts gty (Guarded x l) =
    Guarded (frh x)
            (map (bimap (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty)
                        (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty))
                 l)
  extMap _ ps1 ps2 frh ffd ffe fty fpa fex gts gty (Plain x e) =
    Plain (frh x) (extMap (Proxy @Expr) ps1 ps2 fex ffd ffe fty fpa frh gts gty e)

instance HasExt Expr where
  type HasXField Expr = 'True
  type HasECon Expr = 'False
  type ContainsData Expr = '[FunDef, TypeSig, FunEq, Type, Pattern, RHS]
  extOf (ELit x _) = x
  extOf (EVar x _) = x
  extOf (ECon x _) = x
  extOf (EList x _) = x
  extOf (ETup x _) = x
  extOf (EApp x _ _) = x
  extOf (EOp x _ _ _) = x
  extOf (EIf x _ _ _) = x
  extOf (ECase x _ _) = x
  extOf (ELet x _ _) = x
  extOf (EError x) = x

  extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty = \case
    ELit x l -> ELit (fex x) l
    EVar x n -> EVar (fex x) n
    ECon x n -> ECon (fex x) n
    EList x es -> EList (fex x) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
    ETup x es -> ETup (fex x) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
    EApp x a es -> EApp (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) (map (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty) es)
    EOp x a n b -> EOp (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) n (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b)
    EIf x a b c -> EIf (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty c)
    ECase x a l -> ECase (fex x) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty a)
                         (map (bimap (extMap (Proxy @Pattern) ps1 ps2 fpa)
                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd ffe fty fpa fex gts gty))
                              l)
    ELet x as b -> ELet (fex x) (map (extMap (Proxy @FunDef) ps1 ps2 ffd ffe fty fpa frh fex gts gty) as) (extMap p ps1 ps2 fex ffd ffe fty fpa frh gts gty b)
    EError x -> EError (fex x)