{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE RankNTypes                #-}

module Type.Cache.TH where

import Prelude
import Data.Monoid
import Language.Haskell.TH
import Control.Lens
import Control.Lens.Utils
import Data.Map (Map)
import Data.Typeable
import qualified Data.Map as Map

helperName s = mkName $ "_cache_helper_" <> nameBase s

targetName n = mkName . \case
    Just s  -> trim s
    Nothing -> nameBase n <> "'"
    where trim = unwords . words

cacheHelper n cmn = do
    let hn    = helperName cn
        cn    = targetName n cmn
        udef  = mkName "undefined"
        proxy = mkName "Proxy"
    r <- reify n
    let bndrs  = r ^. tyVarBndrs
        vnames = fmap (VarT . view name) bndrs
    return [ValD (VarP hn) (NormalB (SigE (VarE udef) (ForallT bndrs [] (AppT (ConT proxy) (appsT (ConT n) vnames))  ))) []]

cacheType n cmn = do
    let hn = helperName cn
        cn = targetName n cmn
    rdef <- reify n
    VarI _ tr _ _ <- reify hn


    let -- helpers
        bndrBases    = (nameBase . view name <$>)
        bindUsedName m (s, sn) = case Map.lookup s m of
            Just n  -> n
            Nothing -> sn

        -- binders defined by the user
        defBndrs     = view tyVarBndrs rdef
        defFreeBndrs = fmap (name %~ capturable) defBndrs
        defBases     = bndrBases defBndrs
        defAssoc     = zip defBases defFreeBndrs

        -- binders of the helper declaration
        (bndrs, t) = case tr of
            ForallT bs _ t -> (bs, t)
            t              -> ([], t)
        (AppT _ t')  = t -- Omit the Proxy variable
        bndrs'       = view tyVarBndrs tr
        bindedBases  = bndrBases bndrs'
        bindedAssoc  = zip bindedBases bndrs'
        bindedMap    = Map.fromList bindedAssoc

        -- merged binders used to keep the user defined layout
        finalBndrs   = bindUsedName bindedMap <$> defAssoc

    return [TySynD cn finalBndrs t']

assertTypesEq t t' = if typeOf t == typeOf t' 
    then return []
    else fail "Assertion failed: Generated type does not match the cached one. Please update the cache."

appsT t ts = foldl AppT t ts



class HasName a where
    name :: Lens' a Name


instance HasName TyVarBndr where
    name = lens get set where
        get = \case
            PlainTV  n   -> n
            KindedTV n _ -> n
        set v n = case v of
            PlainTV  _   -> PlainTV n
            KindedTV _ k -> KindedTV n k

class MayHaveTyVarBndrs a where
    tryTyVarBndrs :: Lens' a (Maybe [TyVarBndr])

instance MayHaveTyVarBndrs Info where
    tryTyVarBndrs = lens get set where
        get = \case
            ClassI   d i     -> d ^. tryTyVarBndrs
            ClassOpI n t p f -> t ^. tryTyVarBndrs
            TyConI   d       -> d ^. tryTyVarBndrs
            FamilyI  d i     -> d ^. tryTyVarBndrs
            DataConI n t p f -> t ^. tryTyVarBndrs
            VarI     n t d f -> t ^. tryTyVarBndrs
            TyVarI   n t     -> t ^. tryTyVarBndrs
            _                -> Nothing

        set v x = case v of
            ClassI   d i     -> ClassI     (d & tryTyVarBndrs .~ x) i
            ClassOpI n t p f -> ClassOpI n (t & tryTyVarBndrs .~ x) p f 
            TyConI   d       -> TyConI     (d & tryTyVarBndrs .~ x)       
            FamilyI  d i     -> FamilyI    (d & tryTyVarBndrs .~ x) i     
            DataConI n t p f -> DataConI n (t & tryTyVarBndrs .~ x) p f 
            VarI     n t d f -> VarI     n (t & tryTyVarBndrs .~ x) d f 
            TyVarI   n t     -> TyVarI   n (t & tryTyVarBndrs .~ x)     
            _                -> v

instance MayHaveTyVarBndrs Type where
    tryTyVarBndrs = lens get set where
        get = \case
            ForallT b c t -> Just b
            _             -> Nothing

        set v (maybeToList -> x) = case v of
            ForallT _ c t -> ForallT x c t
            _             -> v

instance MayHaveTyVarBndrs Dec where
    tryTyVarBndrs = lens get set where
        get = \case
            DataD             _ _ x _ _ -> Just x
            NewtypeD          _ _ x _ _ -> Just x
            TySynD              _ x _   -> Just x
            ClassD            _ _ x _ _ -> Just x
            FamilyD           _ _ x _   -> Just x
            ClosedTypeFamilyD   _ x _ _ -> Just x
            _                           -> Nothing 

        set v (maybeToList -> x) = case v of
            DataD             a b _ c d -> DataD             a b x c d
            NewtypeD          a b _ c d -> NewtypeD          a b x c d
            TySynD              b _ c   -> TySynD              b x c  
            ClassD            a b _ c d -> ClassD            a b x c d
            FamilyD           a b _ c   -> FamilyD           a b x c  
            ClosedTypeFamilyD   b _ c d -> ClosedTypeFamilyD   b x c d
            a                           -> a

tyVarBndrs' :: MayHaveTyVarBndrs a => Lens' a [TyVarBndr]
tyVarBndrs' = tryTyVarBndrs . fromMaybeLens (error "Type variables not found")

tyVarBndrs :: MayHaveTyVarBndrs a => Lens' a [TyVarBndr]
tyVarBndrs = tryTyVarBndrs . maybeToListLens

fromMaybeLens :: a -> Lens' (Maybe a) a
fromMaybeLens e = lens (fromMaybe e) (const Just)

maybeToListLens :: Lens' (Maybe [a]) [a]
maybeToListLens = lens maybeToList (const Just)

fromMaybe e = \case
    Just a  -> a
    Nothing -> e

maybeToList = \case
    Just s  -> s
    Nothing -> []

capturable :: Name -> Name
capturable = mkName . nameBase