{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections     #-}

-- | This module is kind of half-assed. I don't have any references and it
-- depends on the inliner.
module Kempe.Monomorphize ( closedModule
                          , MonoM
                          , runMonoM
                          , flattenModule
                          , tryMono
                          , ConsAnn (..)
                          -- * Benchmark
                          , closure
                          , mkModuleMap
                          ) where

import           Control.Arrow              ((&&&))
import           Control.Monad              ((<=<))
import           Control.Monad.Except       (MonadError, throwError)
import           Control.Monad.State.Strict (StateT, gets, runStateT)
import           Data.Bifunctor             (second)
import           Data.Containers.ListUtils  (nubOrd)
import           Data.Foldable              (traverse_)
import           Data.Function              (on)
import           Data.Functor               (($>))
import qualified Data.IntMap                as IM
import           Data.List                  (elemIndex, find, groupBy, partition)
import qualified Data.Map                   as M
import           Data.Maybe                 (fromMaybe, mapMaybe)
import           Data.Semigroup             ((<>))
import qualified Data.Set                   as S
import qualified Data.Text                  as T
import           Data.Tuple                 (swap)
import           Data.Tuple.Ext             (fst3, snd3, thd3)
import           Kempe.AST
import           Kempe.AST.Size
import           Kempe.Error
import           Kempe.Name
import           Kempe.Unique
import           Lens.Micro                 (Lens')
import           Lens.Micro.Mtl             (modifying)
import           Prettyprinter              (Doc, Pretty, vsep)
import           Prettyprinter.Debug

-- | New function names, keyed by name + specialized type
--
-- also max state threaded through.
data RenameEnv = RenameEnv { maxState :: Int
                           , fnEnv    :: M.Map (Unique, StackType ()) Unique
                           , consEnv  :: M.Map (Unique, StackType ()) (Unique, ConsAnn MonoStackType)
                           , szEnv    :: SizeEnv
                           }

type MonoM = StateT RenameEnv (Either (Error ()))

prettyDumpBinds :: (Pretty b, Pretty k) => M.Map k b -> Doc a
prettyDumpBinds b = vsep (prettyBind <$> M.toList b)

maxStateLens :: Lens' RenameEnv Int
maxStateLens f s = fmap (\x -> s { maxState = x }) (f (maxState s))

consEnvLens :: Lens' RenameEnv (M.Map (Unique, StackType ()) (Unique, ConsAnn MonoStackType))
consEnvLens f s = fmap (\x -> s { consEnv = x }) (f (consEnv s))

fnEnvLens :: Lens' RenameEnv (M.Map (Unique, StackType ()) Unique)
fnEnvLens f s = fmap (\x -> s { fnEnv = x }) (f (fnEnv s))

szEnvLens :: Lens' RenameEnv SizeEnv
szEnvLens f s = fmap (\x -> s { szEnv = x }) (f (szEnv s))

runMonoM :: Int -> MonoM a -> Either (Error ()) (a, (Int, SizeEnv))
runMonoM maxI = fmap (second (maxState &&& szEnv)) . flip runStateT (RenameEnv maxI mempty mempty mempty)

freshName :: T.Text -> a -> MonoM (Name a)
freshName n ty = do
    pSt <- gets maxState
    Name n (Unique $ pSt + 1) ty
        <$ modifying maxStateLens (+1)

tryMono :: MonadError (Error ()) m => StackType () -> m MonoStackType
tryMono (StackType is os) | S.null (freeVars (is ++ os)) = pure (is, os)
                          | otherwise = throwError $ MonoFailed ()

-- | A 'ModuleMap' is a map which retrives the 'KempeDecl' associated with
-- a given 'Name'
type ModuleMap a c b = IM.IntMap (KempeDecl a c b)

mkModuleMap :: Declarations a c b -> ModuleMap a c b
mkModuleMap = IM.fromList . concatMap toInt where
    toInt d@(FunDecl _ (Name _ (Unique i) _) _ _ _)   = [(i, d)]
    toInt d@(ExtFnDecl _ (Name _ (Unique i) _) _ _ _) = [(i, d)]
    toInt d@(TyDecl _ _ _ ds)                         =
        let us = unUnique . unique . fst <$> ds
            in (, d) <$> us
    toInt _                                           = []

squishTypeName :: BuiltinTy -> T.Text
squishTypeName TyInt  = "int"
squishTypeName TyBool = "bool"
squishTypeName TyWord = "word"
squishTypeName TyInt8 = "int8"

squishType :: KempeTy a -> T.Text
squishType (TyBuiltin _ b)          = squishTypeName b
squishType (TyNamed _ (Name t _ _)) = T.toLower t
squishType TyVar{}                  = error "not meant to be monomorphized!"
squishType (TyApp _ ty ty')         = squishType ty <> squishType ty'

squishMonoStackType :: MonoStackType -> T.Text
squishMonoStackType (is, os) = foldMap squishType is <> "TT" <> foldMap squishType os

renamePattern :: Pattern (StackType ()) (StackType ()) -> MonoM (Pattern (ConsAnn MonoStackType) (StackType ()))
renamePattern (PatternInt ty i)    = pure $ PatternInt ty i
renamePattern (PatternWildcard ty) = pure $ PatternWildcard ty
renamePattern (PatternBool ty b)   = pure $ PatternBool ty b
renamePattern (PatternCons ty (Name t u _)) = do
    cSt <- gets consEnv
    let (u', ann) = M.findWithDefault (error "Internal error? unfound constructor") (u, flipStackType ty) cSt
        ann' = swap <$> ann
    pure $ PatternCons ann' (Name t u' ann')

renameCase :: (Pattern (StackType ()) (StackType ()), [Atom (StackType ()) (StackType ())]) -> MonoM (Pattern (ConsAnn MonoStackType) (StackType ()), [Atom (ConsAnn MonoStackType) (StackType ())])
renameCase (p, as) = (,) <$> renamePattern p <*> traverse renameAtom as

renameAtom :: Atom (StackType ()) (StackType ()) -> MonoM (Atom (ConsAnn MonoStackType) (StackType ()))
renameAtom (AtBuiltin ty b)         = pure $ AtBuiltin ty b
renameAtom (If ty as as')           = If ty <$> traverse renameAtom as <*> traverse renameAtom as'
renameAtom (IntLit ty i)            = pure $ IntLit ty i
renameAtom (Int8Lit ty i)           = pure $ Int8Lit ty i
renameAtom (WordLit ty w)           = pure $ WordLit ty w
renameAtom (BoolLit ty b)           = pure $ BoolLit ty b
renameAtom (Dip ty as)              = Dip ty <$> traverse renameAtom as
renameAtom (AtName ty (Name t u l)) = do
    mSt <- gets fnEnv
    let u' = M.findWithDefault u (u, ty) mSt
    pure $ AtName ty (Name t u' l)
renameAtom (Case ty ls)             = Case ty <$> traverse renameCase ls
renameAtom (AtCons ty (Name t u _)) = do
    cSt <- gets consEnv
    let (u', ann) = M.findWithDefault (error "Internal error? unfound constructor") (u, ty) cSt
    pure $ AtCons ann (Name t u' ann)

renameDecl :: KempeDecl () (StackType ()) (StackType ()) -> MonoM (KempeDecl () (ConsAnn MonoStackType) (StackType ()))
renameDecl (FunDecl l n is os as) = FunDecl l n is os <$> traverse renameAtom as
renameDecl (Export ty abi (Name t u l)) = do
    mSt <- gets fnEnv
    let u' = M.findWithDefault (error "Shouldn't happen; might be user error or internal error") (u, ty) mSt
    pure $ Export ty abi (Name t u' l)
renameDecl (ExtFnDecl l n tys tys' b) = pure $ ExtFnDecl l n tys tys' b
renameDecl (TyDecl l n vars ls)       = pure $ TyDecl l n vars ls

-- | Call 'closedModule' and perform any necessary renamings
flattenModule :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (ConsAnn MonoStackType) (StackType ()))
flattenModule = renameMonoM <=< closedModule

-- | To be called after 'closedModule'
renameMonoM :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (ConsAnn MonoStackType) (StackType ()))
renameMonoM = traverse renameDecl

-- | Filter so that only the 'KempeDecl's necessary for exports are there, and
-- fan out top-level functions into all necessary specializations.
--
-- This will throw an exception on ill-typed programs.
--
-- The 'Module' returned will have to be renamed.
closedModule :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (StackType ()) (StackType ()))
closedModule m = addExports <$> do
    { fn' <- traverse (uncurry specializeDecl . drop1) fnDecls
    ; traverse_ insTyDecl $ nubOrd (snd3 <$> tyDecls)
    ; ty' <- specializeTyDecls tyDecls
    ; pure (ty' ++ fn')
    }
    where addExports = (++ exportsOnly m)
          key = mkModuleMap m
          roots = S.toList $ closure (m, key)
          gatherDecl (n@(Name _ (Unique i) _), ty) = -- TODO: findWithDefault?
            case IM.lookup i key of
                Just decl -> (n, decl, ty)
                Nothing   -> error "Internal error! module map should contain all names."
          rootDecl = gatherDecl <$> roots -- FIXME: two-steps away, the roots are not monomorphized! So it tries to create specialized declarations of type a b -- a b a &c.
          drop1 ~(_, y, z) = (y, z)
          (tyDecls, fnDecls) = partition (isTyDecl . snd3) rootDecl
          isTyDecl TyDecl{} = True
          isTyDecl _        = False

-- group specializations by type name?
specializeTyDecls :: [(TyName (StackType ()), KempeDecl () (StackType ()) (StackType ()), StackType ())] -> MonoM [KempeDecl () (StackType ()) (StackType ())]
specializeTyDecls ds = traverse (uncurry mkTyDecl) processed
    where toMerge = groupBy ((==) `on` snd3) ds
          processed = fmap process toMerge
          process tyDs@((_, x, _):_) = (x, zip (fst3 <$> tyDs) (thd3 <$> tyDs))
          process []                 = error "Empty group!"

isTyVar :: KempeTy a -> Bool
isTyVar TyVar{} = True
isTyVar _       = False

extrNames :: KempeTy a -> Name a
extrNames (TyVar _ n) = n
extrNames _           = error "Internal error!"

sizeLeaf :: [Name a] -- ^ Type variables as declared
         -> [KempeTy a]
         -> MonoM Size
sizeLeaf fv tys = do
    { let (tvs, conc) = partition isTyVar tys
    ; pad <- sizeStack <$> gets szEnv <*> pure conc
    ; let tvPrecompose = fmap (forVar . extrNames) tvs
    ; let tvComposed = foldr compose (const pad) tvPrecompose
    ; pure tvComposed
    }
  where
    findIx x = fromMaybe (error "Internal error: can't find index of type variable.") $ elemIndex x fv
    forVar n =
        let i = findIx n
            in (!! i)
    compose sz sz' = \tys' -> sz tys' + sz' tys'

insTyDecl :: KempeDecl a c b -> MonoM ()
insTyDecl (TyDecl _ (Name _ (Unique k) _) fv leaves) = do
    leafSizes <- traverse (sizeLeaf fv) (fmap snd leaves)
    let consSz = \tys -> 1 + maximum (($tys) <$> leafSizes) -- for the tag
    modifying szEnvLens (IM.insert k consSz)
insTyDecl _ = error "Shouldn't happen."

mkTyDecl :: KempeDecl () (StackType ()) (StackType ()) -> [(TyName (StackType ()), StackType ())] -> MonoM (KempeDecl () (StackType ()) (StackType ()))
mkTyDecl (TyDecl _ tn ns preConstrs) constrs = do
    env <- gets szEnv
    renCons <- traverse (\(tn', ty) -> do { ty'@(is, _) <- tryMono ty ; (, is) <$> renamedCons (tn' $> ty') ty' (ConsAnn (szType env ty') (getTag tn')) }) constrs
    pure $ TyDecl () tn ns renCons
    where indexAt p xs = fst $ fromMaybe (error "Internal error.") $ find (\(_, x) -> p x) (zip [0..] xs)
          getTag (Name _ u _) = indexAt (== u) preIxes
          preIxes = fmap (unique . fst) preConstrs
          szType env (_, [o]) = size' env o
          szType _ _          = error "Internal error: ill-typed constructor."
mkTyDecl _ _ = error "Shouldn't happen."

specializeDecl :: KempeDecl () (StackType ()) (StackType ()) -> StackType () -> MonoM (KempeDecl () (StackType ()) (StackType ()))
specializeDecl (FunDecl _ n _ _ as) sty = do
    (Name t u newStackType@(StackType is os)) <- renamed n =<< tryMono sty
    pure $ FunDecl newStackType (Name t u newStackType) is os as
specializeDecl (ExtFnDecl l n tys tys' b) _ = pure $ ExtFnDecl l n tys tys' b
specializeDecl (Export l abi n) _           = pure $ Export l abi n
specializeDecl TyDecl{} _                   = error "Shouldn't happen."
-- leave exports and foreign imports alone (have to be monomorphic)

renamedCons :: TyName a -> MonoStackType -> (MonoStackType -> ConsAnn MonoStackType) -> MonoM (TyName (StackType ()))
renamedCons (Name t i _) sty@(is, os) fAnn = do
    let t' = t <> squishMonoStackType sty
    (Name _ j _) <- freshName t' sty
    let newStackType = StackType is os
        ann = fAnn sty
    modifying consEnvLens (M.insert (i, newStackType) (j, ann))
    pure (Name t' j newStackType)

-- | Insert a specialized rename.
renamed :: Name a -> MonoStackType -> MonoM (Name (StackType ()))
renamed (Name t i _) sty@(is, os) = do
    let t' = t <> squishMonoStackType sty
    (Name _ j _) <- freshName t' sty
    let newStackType = StackType is os
    modifying fnEnvLens (M.insert (i, newStackType) j)
    pure (Name t' j newStackType)

closure :: (Declarations a (StackType ()) (StackType ()), ModuleMap a (StackType ()) (StackType ())) -> S.Set (Name (StackType ()), StackType ())
closure (m, key) = loop roots S.empty
    where roots = S.fromList (exports m)
          loop ns avoid =
            let res = foldMap (step . fst) (ns S.\\ avoid)
                in if res == ns
                    then res
                    else ns <> loop res (ns <> avoid)
          step (Name _ (Unique i) _) =
            case IM.lookup i key of
                Just decl -> namesInDecl decl
                Nothing   -> error "Internal error! module map should contain all names."

namesInDecl :: KempeDecl a (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
namesInDecl TyDecl{}             = S.empty
namesInDecl ExtFnDecl{}          = S.empty
namesInDecl Export{}             = S.empty
namesInDecl (FunDecl _ _ _ _ as) = foldMap namesInAtom as

namesInAtom :: Atom (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
namesInAtom AtBuiltin{}                = S.empty
namesInAtom (If _ as as')              = foldMap namesInAtom as <> foldMap namesInAtom as'
namesInAtom (Dip _ as)                 = foldMap namesInAtom as
namesInAtom (AtName _ n@(Name _ _ l))  = S.singleton (n, l)
namesInAtom (AtCons _ tn@(Name _ _ l)) = S.singleton (tn, l)
namesInAtom IntLit{}                   = S.empty
namesInAtom BoolLit{}                  = S.empty
namesInAtom Int8Lit{}                  = S.empty
namesInAtom WordLit{}                  = S.empty
namesInAtom (Case _ as)                = foldMap namesInAtom (foldMap snd as) <> foldMap (namesInPattern . fst) as

namesInPattern :: Pattern (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
namesInPattern (PatternCons _ tn@(Name _ _ l)) = S.singleton (tn, flipStackType l) -- we have to flipStackType here because the type is the reverse of the constructor that we want to look up
namesInPattern _                               = S.empty

exports :: Declarations a c b -> [(Name b, b)]
exports = mapMaybe exportsDecl

exportsOnly :: Declarations a c b -> Declarations a c b
exportsOnly = mapMaybe getExport where
    getExport d@Export{} = Just d
    getExport _          = Nothing

exportsDecl :: KempeDecl a c b -> Maybe (Name b, b)
exportsDecl (Export _ _ n@(Name _ _ l)) = Just (n, l)
exportsDecl _                           = Nothing