{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE LambdaCase #-}
module HSVIS.AST where

import Data.Bifunctor (bimap, second)
import qualified Data.Kind as DK
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 :: f s -> X f s
  extMap :: Proxy f -> Proxy s1 -> Proxy s2
         -> XMapperTelescope (TransitiveContainsData '[f]) s1 s2
              (EMapperTelescope (TransitiveContainsData '[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 ASTShowPrereqs f s :: DK.Constraint where
  ASTShowPrereqs f s =
    (If (HasXField f) (Show (X f s)) (() :: DK.Constraint)
    ,If (HasECon f) (Show (E f s)) (() :: DK.Constraint)
    ,ASTShowPrereqsC (ContainsData f) s)

type family ASTShowPrereqsC l s :: DK.Constraint where
  ASTShowPrereqsC '[] s = ()
  ASTShowPrereqsC (f ': fs) s = (Show (f s), ASTShowPrereqsC fs s)

type family a && b where
  'False && b = 'False
  'True && b = b

type family In x l where
  In x '[] = 'False
  In x (x ': ys) = 'True
  In x (_ ': ys) = In x ys

type family Sublist l1 l2 where
  Sublist '[] l = 'True
  Sublist (x ': xs) l = In x l && Sublist xs l

type family Union l1 l2 where
  Union '[] l = l
  Union (x ': xs) l = If (In x l) (Union xs l) (x ': Union xs l)

type family MapContainsData l where
  MapContainsData '[] = '[]
  MapContainsData (f ': fs) = Union (ContainsData f) (MapContainsData fs)

type family TransitiveContainsData fs where
  TransitiveContainsData fs = TransitiveContainsData' fs (Sublist (Union fs (MapContainsData fs)) fs)

type family TransitiveContainsData' fs isfull where
  TransitiveContainsData' fs 'True = fs
  TransitiveContainsData' fs 'False = TransitiveContainsData (Union fs (MapContainsData fs))

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)
       ((X f s1 -> E f s1 -> f s2) -> EMapperTelescope fs s1 s2 a)
       (EMapperTelescope fs s1 s2 a)


newtype Name = Name String
  deriving (Show, Eq)

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 [Name] [(Name, [Type s])]
deriving instance ASTShowPrereqs DataDef s => Show (DataDef s)

data FunDef s = FunDef (X FunDef s) Name (Maybe (Type s)) (NonEmpty (FunEq s))
deriving instance ASTShowPrereqs FunDef s => Show (FunDef s)

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

data Kind s
  = KType (X Kind s)
  | KFun (X Kind s) (Kind s) (Kind s)
deriving instance ASTShowPrereqs Kind s => Show (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

  -- extension point
  | TExt (X Type s) !(E Type s)
deriving instance ASTShowPrereqs Type s => Show (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 ASTShowPrereqs 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 ASTShowPrereqs RHS 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)]
  | ELet  (X Expr s) [FunDef s] (Expr s)
  | EError (X Expr s)
deriving instance ASTShowPrereqs Expr 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 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 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 = '[Type, FunEq]
  extOf (FunDef x _ _ _) = x
  extMap _ ps1 ps2 ffd ffe frh fty fpa fex gty (FunDef x n msig es) =
    FunDef (ffd x) n (extMap (Proxy @Type) ps1 ps2 fty gty <$> msig)
           (fmap (extMap (Proxy @FunEq) ps1 ps2 frh ffd fpa fex fty ffe gty) es)

instance HasExt FunEq where
  type HasXField FunEq = 'True
  type HasECon FunEq = 'False
  type ContainsData FunEq = '[Pattern, RHS]
  extOf (FunEq x _ _ _) = x
  extMap _ ps1 ps2 frh ffd fpa fex fty ffe 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 fpa fex fty ffe gty rhs)

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

instance HasExt Type where
  type HasXField Type = 'True
  type HasECon Type = 'True
  type ContainsData Type = '[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 (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 _ _   _   _ g (TExt x e) = g x e

instance HasExt Pattern where
  type HasXField Pattern = 'True
  type HasECon Pattern = 'False
  type ContainsData Pattern = '[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 = '[Expr]
  extOf (Guarded x _) = x
  extOf (Plain x _) = x

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

instance HasExt Expr where
  type HasXField Expr = 'True
  type HasECon Expr = 'False
  type ContainsData Expr = '[Expr, Pattern, RHS, FunDef]
  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 frh ffd fpa fex fty ffe 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 frh ffd fpa fex fty ffe gty) es)
    ETup x es -> ETup (fex x) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
    EApp x a es -> EApp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
    EOp x a n b -> EOp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) n (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
    EIf x a b c -> EIf (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty c)
    ECase x a l -> ECase (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a)
                         (map (bimap (extMap (Proxy @Pattern) ps1 ps2 fpa)
                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe gty))
                              l)
    ELet x as b -> ELet (fex x) (map (extMap (Proxy @FunDef) ps1 ps2 ffd ffe frh fty fpa fex gty) as) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
    EError x -> EError (fex x)