{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Frontend AST
module Kempe.AST ( BuiltinTy (..)
                 , KempeTy (..)
                 , StackType (..)
                 , ConsAnn (..)
                 , Atom (..)
                 , BuiltinFn (..)
                 , KempeDecl (..)
                 , Pattern (..)
                 , ABI (..)
                 , Module
                 , freeVars
                 , MonoStackType
                 , SizeEnv
                 , size
                 , sizeStack
                 , prettyMonoStackType
                 , prettyTyped
                 , prettyTypedModule
                 , prettyFancyModule
                 , prettyModule
                 , flipStackType
                 -- * I resent this...
                 , voidStackType
                 ) where

import           Control.DeepSeq         (NFData)
import           Data.Bifunctor          (Bifunctor (..))
import qualified Data.ByteString.Lazy    as BSL
import           Data.Foldable           (toList)
import           Data.Functor            (void)
import           Data.Int                (Int64, Int8)
import qualified Data.IntMap             as IM
import           Data.List.NonEmpty      (NonEmpty)
import qualified Data.List.NonEmpty      as NE
import           Data.Monoid             (Sum (..))
import qualified Data.Set                as S
import           Data.Text.Lazy.Encoding (decodeUtf8)
import           Data.Word               (Word8)
import           GHC.Generics            (Generic)
import           Kempe.Name
import           Kempe.Unique
import           Numeric.Natural
import           Prettyprinter           (Doc, Pretty (pretty), align, braces, brackets, colon, concatWith, fillSep, hsep, parens, pipe, sep, vsep, (<+>))
import           Prettyprinter.Ext

data BuiltinTy = TyInt
               | TyBool
               | TyInt8
               | TyWord
               deriving (Generic, NFData, Eq, Ord)

instance Pretty BuiltinTy where
    pretty TyInt  = "Int"
    pretty TyBool = "Bool"
    pretty TyInt8 = "Int8"
    pretty TyWord = "Word"

-- equality for sum types &c.

data KempeTy a = TyBuiltin a BuiltinTy
               | TyNamed a (TyName a)
               | TyVar a (Name a)
               | TyApp a (KempeTy a) (KempeTy a) -- type applied to another, e.g. Just Int
               deriving (Generic, NFData, Functor, Eq, Ord) -- questionable eq instance but eh

data StackType b = StackType { quantify :: S.Set (Name b)
                             , inTypes  :: [KempeTy b]
                             , outTypes :: [KempeTy b]
                             } deriving (Generic, NFData, Eq, Ord)

type MonoStackType = ([KempeTy ()], [KempeTy ()])

-- | Annotation carried on constructors to keep size information through the IR
-- generation phase.
data ConsAnn a = ConsAnn { tySz :: Int64, tag :: Word8, consTy :: a }
    deriving (Functor, Foldable, Traversable, Generic, NFData)

instance Pretty a => Pretty (ConsAnn a) where
    pretty (ConsAnn tSz b ty) = braces ("tySz" <+> colon <+> pretty tSz <+> "tag" <+> colon <+> pretty b <+> "type" <+> colon <+> pretty ty)

prettyMonoStackType :: MonoStackType -> Doc a
prettyMonoStackType (is, os) = sep (fmap pretty is) <+> "--" <+> sep (fmap pretty os)

instance Pretty (StackType a) where
    pretty (StackType _ ins outs) = sep (fmap pretty ins) <+> "--" <+> sep (fmap pretty outs)

voidStackType :: StackType a -> StackType ()
voidStackType (StackType vars ins outs) = StackType (S.map void vars) (void <$> ins) (void <$> outs)

instance Pretty (KempeTy a) where
    pretty (TyBuiltin _ b)  = pretty b
    pretty (TyNamed _ tn)   = pretty tn
    pretty (TyVar _ n)      = pretty n
    pretty (TyApp _ ty ty') = parens (pretty ty <+> pretty ty')

data Pattern c b = PatternInt b Integer
                 | PatternCons { patternKind :: c, patternName :: TyName c } -- a constructed pattern
                 | PatternWildcard b
                 | PatternBool b Bool
                 deriving (Eq, Ord, Generic, NFData, Functor, Foldable, Traversable)

instance Bifunctor Pattern where
    second = fmap
    first f (PatternCons l tn)  = PatternCons (f l) (fmap f tn)
    first _ (PatternInt l i)    = PatternInt l i
    first _ (PatternWildcard l) = PatternWildcard l
    first _ (PatternBool l b)   = PatternBool l b

instance Pretty (Pattern c a) where
    pretty (PatternInt _ i)   = pretty i
    pretty (PatternBool _ b)  = pretty b
    pretty PatternWildcard{}  = "_"
    pretty (PatternCons _ tn) = pretty tn

prettyTypedPattern :: Pattern (StackType ()) (StackType ()) -> Doc ann
prettyTypedPattern (PatternCons ty tn) = parens (pretty tn <+> ":" <+> pretty ty)
prettyTypedPattern p                   = pretty p

instance Pretty (Atom c a) where
    pretty (AtName _ n)    = pretty n
    pretty (Dip _ as)      = "dip(" <> fillSep (fmap pretty as) <> ")"
    pretty (AtBuiltin _ b) = pretty b
    pretty (AtCons _ tn)   = pretty tn
    pretty (If _ as as')   = "if(" <> align (fillSep (fmap pretty as)) <> ", " <> align (fillSep (fmap pretty as')) <> ")"
    pretty (IntLit _ i)    = pretty i
    pretty (BoolLit _ b)   = pretty b
    pretty (WordLit _ w)   = pretty w <> "u"
    pretty (Int8Lit _ i)   = pretty i <> "i8"
    pretty (Case _ ls)     = "case" <+> braces (align (vsep (toList $ fmap (uncurry prettyLeaf) ls)))

prettyLeaf :: Pattern c a -> [Atom c a] -> Doc ann
prettyLeaf p as = pipe <+> pretty p <+> "->" <+> align (fillSep (fmap pretty as))

prettyTypedLeaf :: Pattern (StackType ()) (StackType ()) -> [Atom (StackType ()) (StackType ())] -> Doc ann
prettyTypedLeaf p as = pipe <+> prettyTypedPattern p <+> "->" <+> align (fillSep (fmap prettyTyped as))

prettyTyped :: Atom (StackType ()) (StackType ()) -> Doc ann
prettyTyped (AtName ty n)    = parens (pretty n <+> ":" <+> pretty ty)
prettyTyped (Dip _ as)       = "dip(" <> fillSep (prettyTyped <$> as) <> ")"
prettyTyped (AtBuiltin ty b) = parens (pretty b <+> ":" <+> pretty ty)
prettyTyped (AtCons ty tn)   = parens (pretty tn <+> ":" <+> pretty ty)
prettyTyped (If _ as as')    = "if(" <> align (fillSep (prettyTyped <$> as)) <> ", " <> align (fillSep (prettyTyped <$> as')) <> ")"
prettyTyped (IntLit _ i)     = pretty i
prettyTyped (BoolLit _ b)    = pretty b
prettyTyped (Int8Lit _ i)    = pretty i <> "i8"
prettyTyped (WordLit _ n)    = pretty n <> "u"
prettyTyped (Case _ ls)      = "case" <+> braces (align (vsep (toList $ fmap (uncurry prettyTypedLeaf) ls)))

data Atom c b = AtName b (Name b)
              | Case b (NonEmpty (Pattern c b, [Atom c b]))
              | If b [Atom c b] [Atom c b]
              | Dip b [Atom c b]
              | IntLit b Integer
              | WordLit b Natural
              | Int8Lit b Int8
              | BoolLit b Bool
              | AtBuiltin b BuiltinFn
              | AtCons c (TyName c)
              deriving (Eq, Ord, Generic, NFData, Functor, Foldable, Traversable)

instance Bifunctor Atom where
    second = fmap
    first f (AtCons l n)    = AtCons (f l) (fmap f n)
    first _ (AtName l n)    = AtName l n
    first _ (IntLit l i)    = IntLit l i
    first _ (WordLit l w)   = WordLit l w
    first _ (Int8Lit l i)   = Int8Lit l i
    first _ (BoolLit l b)   = BoolLit l b
    first _ (AtBuiltin l b) = AtBuiltin l b
    first f (Dip l as)      = Dip l (fmap (first f) as)
    first f (If l as as')   = If l (fmap (first f) as) (fmap (first f) as')
    first f (Case l ls)     =
        let (ps, aLs) = NE.unzip ls
            in Case l $ NE.zip (fmap (first f) ps) (fmap (fmap (first f)) aLs)

data BuiltinFn = Drop
               | Swap
               | Dup
               | IntPlus
               | IntMinus
               | IntTimes
               | IntDiv
               | IntMod
               | IntEq
               | IntLeq
               | IntLt
               | IntGeq
               | IntGt
               | IntNeq
               | IntShiftR
               | IntShiftL
               | IntXor
               | WordPlus
               | WordTimes
               | WordMinus
               | WordDiv
               | WordMod
               | WordShiftR
               | WordShiftL
               | WordXor
               | And
               | Or
               | Xor
               | IntNeg
               | Popcount
               deriving (Eq, Ord, Generic, NFData)

instance Pretty BuiltinFn where
    pretty Drop       = "drop"
    pretty Swap       = "swap"
    pretty Dup        = "dup"
    pretty IntPlus    = "+"
    pretty IntMinus   = "-"
    pretty IntTimes   = "*"
    pretty IntDiv     = "/"
    pretty IntMod     = "%"
    pretty IntEq      = "="
    pretty IntLeq     = "<="
    pretty IntLt      = "<"
    pretty IntShiftR  = ">>"
    pretty IntShiftL  = "<<"
    pretty WordPlus   = "+~"
    pretty WordTimes  = "*~"
    pretty WordShiftL = "<<~"
    pretty WordShiftR = ">>~"
    pretty IntXor     = "xori"
    pretty WordXor    = "xoru"
    pretty IntGeq     = ">="
    pretty IntGt      = ">"
    pretty IntNeq     = "!="
    pretty WordMinus  = "-~"
    pretty WordDiv    = "/~"
    pretty WordMod    = "%~"
    pretty And        = "&"
    pretty Or         = "||"
    pretty Xor        = "xor"
    pretty IntNeg     = "~"
    pretty Popcount   = "popcount"

data ABI = Cabi
         | Kabi
         deriving (Eq, Ord, Generic, NFData)

instance Pretty ABI where
    pretty Cabi = "cabi"
    pretty Kabi = "kabi"

prettyKempeDecl :: (Atom c b -> Doc ann) -> KempeDecl a c b -> Doc ann
prettyKempeDecl atomizer (FunDecl _ n is os as) = pretty n <+> align (":" <+> sep (fmap pretty is) <+> "--" <+> sep (fmap pretty os) <#> "=:" <+> brackets (align (fillSep (atomizer <$> as))))
prettyKempeDecl _ (Export _ abi n)              = "%foreign" <+> pretty abi <+> pretty n
prettyKempeDecl _ (ExtFnDecl _ n is os b)       = pretty n <+> align (":" <+> sep (fmap pretty is) <+> "--" <+> sep (fmap pretty os) <#> "=:" <+> "$cfun" <> pretty (decodeUtf8 b))
prettyKempeDecl _ (TyDecl _ tn ns ls)           = "type" <+> pretty tn <+> hsep (fmap pretty ns) <+> braces (concatWith (\x y -> x <+> pipe <+> y) $ fmap (uncurry prettyTyLeaf) ls)

instance Pretty (KempeDecl a b c) where
    pretty = prettyKempeDecl pretty

prettyTyLeaf :: TyName a -> [KempeTy b] -> Doc ann
prettyTyLeaf cn vars = pretty cn <+> hsep (fmap pretty vars)

-- TODO: separate annotations for TyName in TyDecl
data KempeDecl a c b = TyDecl a (TyName a) [Name a] [(TyName b, [KempeTy a])]
                     | FunDecl b (Name b) [KempeTy a] [KempeTy a] [Atom c b]
                     | ExtFnDecl b (Name b) [KempeTy a] [KempeTy a] BSL.ByteString -- ShortByteString?
                     | Export b ABI (Name b)
                     deriving (Eq, Ord, Generic, NFData, Functor, Foldable, Traversable)

instance Bifunctor (KempeDecl a) where
    first _ (TyDecl x tn ns ls)        = TyDecl x tn ns ls
    first f (FunDecl l n tys tys' as)  = FunDecl l n tys tys' (fmap (first f) as)
    first _ (ExtFnDecl l n tys tys' b) = ExtFnDecl l n tys tys' b
    first _ (Export l abi n)           = Export l abi n
    second = fmap

prettyModuleGeneral :: (Atom c b -> Doc ann) -> Module a c b -> Doc ann
prettyModuleGeneral atomizer = sep . fmap (prettyKempeDecl atomizer)

prettyFancyModule :: Module () (ConsAnn (StackType ())) (StackType ()) -> Doc ann
prettyFancyModule = prettyTypedModule . fmap (first consTy)

prettyTypedModule :: Module () (StackType ()) (StackType ()) -> Doc ann
prettyTypedModule = prettyModuleGeneral prettyTyped

prettyModule :: Module a c b -> Doc ann
prettyModule = prettyModuleGeneral pretty

type Module a c b = [KempeDecl a c b]

extrVars :: KempeTy a -> [Name a]
extrVars TyBuiltin{}      = []
extrVars TyNamed{}        = []
extrVars (TyVar _ n)      = [n]
extrVars (TyApp _ ty ty') = extrVars ty ++ extrVars ty'

freeVars :: [KempeTy a] -> S.Set (Name a)
freeVars tys = S.fromList (concatMap extrVars tys)

type SizeEnv = IM.IntMap Int64

-- the kempe sizing system is kind of fucked (it works tho)

-- | Don't call this on ill-kinded types; it won't throw any error.
size :: SizeEnv -> KempeTy a -> Int64
size _ (TyBuiltin _ TyInt)                 = 8 -- since we're only targeting x86_64 and aarch64 we have 64-bit 'Int's
size _ (TyBuiltin _ TyBool)                = 1
size _ (TyBuiltin _ TyInt8)                = 1
size _ (TyBuiltin _ TyWord)                = 8
size _ TyVar{}                             = error "Internal error: type variables should not be present at this stage."
size env (TyNamed _ (Name _ (Unique k) _)) = IM.findWithDefault (error "Size not in map!") k env
size env (TyApp _ ty ty')                  = size env ty + size env ty'

sizeStack :: SizeEnv -> [KempeTy a] -> Int64
sizeStack env = getSum . foldMap (Sum . size env)

-- | Used in "Kempe.Monomorphize" for patterns
flipStackType :: StackType () -> StackType ()
flipStackType (StackType vars is os) = StackType vars os is