{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS -Wall #-}
module DataTreeView.DataToTree where

import Control.Applicative
import Control.Exception.Lifted
import Control.Monad.Base
import Control.Monad.Reader
import Control.Monad.Trans.Control
import Data.Data
import Data.List(intercalate)
import Data.Monoid
import Data.String.Utils(replace)
import DataTreeView.Row
import DataTreeView.StrictTypes
import Prelude hiding(catch)
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Lazy.Char8 as BL
import qualified Data.IntMap as IntMap
import qualified Data.Map as Map
import qualified Data.Set as Set





-- | Functionality a @CustomHandler@ may use.
data CustomHandlerServices = CHS {
    chsSelf     :: forall d. Data d     => d -> MCH (StrictTree Row) 
}
    deriving(Typeable)



newtype MCH a = MCH (ReaderT CustomHandlerServices IO a) 
    deriving(Functor, Applicative, Monad, MonadIO, MonadBase IO, Typeable) 

runMCH ::  MCH a -> CustomHandlerServices -> IO a
runMCH (MCH x) = runReaderT x

instance MonadBaseControl IO MCH where
    data StM MCH a = Stm_MCH { unStm_MCH :: StM (ReaderT CustomHandlerServices IO) a }

    liftBaseWith (f :: RunInBase MCH IO -> IO a) = MCH (liftBaseWith f')
        where
            f' :: RunInBase (ReaderT CustomHandlerServices IO) IO -> IO a
            f' rib' = f rib

                where
                    rib :: forall a'. MCH a' -> IO (StM MCH a')
                    rib = liftM Stm_MCH . rib' . (\(MCH x) -> x)

    restoreM = MCH . restoreM . unStm_MCH


-- | Invokes the /final/ Data-to-Tree conversion function, which includes the generic handler, the 'CustomHandler' being defined, and any 'CustomHandler's 'mappend'ed to the one being defined.
-- 
-- Thus, invoking @'self' x@ from your implementation of @'runCH' x@ will usually amount to an infinite loop, but invoking @'self' y@ on some child @y@ of @x@ is fine.
self ::  Data d => d -> MCH (StrictTree Row)
self x = do
    f <- MCH (asks chsSelf)
    f x






-- | Custom handlers are used for overriding the generic formatting behaviour at some values (for example, printing a @String@ directly into the /Constructor or value/ column rather than as a list of chars). 
--
-- If a custom handler matches (returns a 'Just'), no recursion into its children is performed by default. You can recurse manually using 'self'. 
newtype CustomHandler = CH { runCH :: forall d. Data d => d -> MCH (Maybe (StrictTree Row)) }


-- | Creates a custom handler that only applies at a fixed monomorphic type @a@.
monoCH :: Typeable a => (a -> MCH (Maybe (StrictTree Row))) -> CustomHandler
monoCH f = CH (case cast f of
                    Nothing -> \_ -> return Nothing
                    Just f' -> f')

-- | Creates a pure custom handler that only applies at a fixed monomorphic type @a@.
monoPureCH :: Typeable a => (a -> (Maybe (StrictTree Row))) -> CustomHandler
monoPureCH f = CH (case cast f of
                     Nothing -> \_ -> return Nothing
                     Just f' -> return . f')

-- | Creates a pure custom handler that only and always applies at a fixed monomorphic type @a@.
monoPureCH' :: Typeable a => (a -> StrictTree Row) -> CustomHandler
monoPureCH' f = CH (case cast f of
                     Nothing -> \_ -> return Nothing
                     Just f' -> return . Just . f')

-- | Creates a custom handler that applies at @f a@ for a fixed type constructor @f@ and for any @a@.
poly1CH :: Typeable1 f => (forall a. Data a => f a -> MCH (Maybe (StrictTree Row))) -> CustomHandler
poly1CH f = CH (case dataCast1 (Q f) of
                     Nothing -> \_ -> return Nothing
                     Just (Q f') -> f')

-- | Creates a custom handler that applies at @f a b@ for a fixed binary type constructor @f@ and for any @a@ and @b@.
poly2CH :: Typeable2 f => (forall a b. (Data a, Data b) => f a b -> MCH (Maybe (StrictTree Row))) -> CustomHandler
poly2CH f = CH (case dataCast2 (Q f) of
                     Nothing -> \_ -> return Nothing
                     Just (Q f') -> f')

-- from syb
newtype Q q x = Q { unQ :: x -> q }


-- | The 'mempty' 'CustomHandler' handles nothing. @'mappend' ch1 ch2@ tries @ch1@ first, falling back to @ch2@ if @ch1@ doesn't handle the argument.
instance Monoid CustomHandler where
    mempty = CH (\_ -> return Nothing)
    mappend ch1 ch2 = CH (\x -> do
        y1 <- runCH ch1 x
        case y1 of
              Just _ -> return y1
              Nothing -> runCH ch2 x)

-- | Construct a 'Row' with empty 'rowFieldName'. 
--
-- You usually can't know the field name for a node @x@ during a recursive call with @x@ at its root - the field name will be added /afterwards/ by the generic handler for the parent node of @x@ if that parent node is a constructor with named fields.
newRow ::      [CellAttr] -- ^ 'rowCV'
            -> [CellAttr] -- ^ 'rowCustomInfo'
            -> [CellAttr] -- ^ 'rowTypeName'
            -> Row
newRow (cellData -> rowCV) (cellData -> rowCustomInfo) (cellData -> rowTypeName) = 
    Row {rowCV,rowTypeName,rowCustomInfo,rowFieldName = mempty}

dataToTree :: forall d. Data d => CustomHandler -> d -> IO (StrictTree Row)
dataToTree ch = \x -> runMCH (chsSelf x) customHandlerServices
    where
        customHandlerServices = CHS {chsSelf} 

        ch' = ch `mappend` builtinCH 

        chsSelf :: forall e. Data e  => e -> MCH (StrictTree Row)
        chsSelf x = chsSelf1 x `catch` excHandler
        
        chsSelf1 :: forall e. Data e  => e -> MCH (StrictTree Row)
        chsSelf1 x = do 
            y <- {-# SCC "dataToTree/runCH" #-} runCH ch' x 
            z <- case y of
                      Just y' -> return y'
                      Nothing -> genericHandler x
            evaluate z

excHandler :: SomeException -> MCH (StrictTree Row)
excHandler e = return (strictTree (formatExc e))



formatExc :: SomeException -> Row
formatExc (SomeException e) = 
    newRow 
        (colored [txt $ show e]) 
        (colored [txt $ showAttrList [("exception type",show (typeOf e))]]) 
        mempty
 where
    colored = ([ bgcolor $ "black", fgcolor $ "white" ] ++)


-- | Generates the subtree using the 'Data' instance of the argument (calls 'self' on the children, not 'genericHandler').
--
-- You can invoke this from your 'CustomHandler' and then override some attributes of the result.
genericHandler :: forall e. Data e => e -> MCH (StrictTree Row)
genericHandler x = do
    rec <- {-# SCC "genericHandler/rec" #-} sequence (gmapQ self x)

    ctorRow <- (do 
                    let constr = toConstr x
                    ctorText <- (evaluate . showConstr) constr
                    return $ case constrRep constr of
                        AlgConstr _ -> newRow [ txt $ ctorText ] [] [txt $ showTypeOf x]
                        _ -> formatLit ctorText "" (typeOf x)
                )
                `catch` (return . formatExc)

        -- Add record field names
    let rec'   =  {-# SCC "genericHandler/rec'" #-}
                  case safeConstrFields x of                                                    
                    Nothing -> rec                                                            
                    Just fieldNames ->                                                        
                     zipWith                                                                  
                      (\n fn -> modifyValue n (addFieldName fn)) 
                        rec                                                                    
                        (fieldNames ++ repeat "")                                              


    return $ strictTree (ctorRow, rec')

safeConstrFields :: Data a => a -> Maybe [String]
safeConstrFields x = 
    case dataTypeRep (dataTypeOf x) of
         AlgRep _ -> case constrFields (toConstr x) of
                          [] -> Nothing
                          xs -> Just xs
         _ -> Nothing

-- safeConstrFields :: Constr -> Maybe [String]
-- safeConstrFields c = unsafePerformIO (
--     (Just <$> evaluate (constrFields c)) `Control.Exception.catch` (\(SomeException e) -> return Nothing) 
--     )


builtinCH :: CustomHandler
builtinCH = mconcat [strCH, lstCH, mapCH, setCH, intMapCH, bsCH, blCH, funCH]

lstCH ::  CustomHandler
lstCH = container1CH (\x -> ("List",[("length",(show . length) x)],fmap AnyData x))

mapCH :: CustomHandler
mapCH = container2CH (\x -> ("Map",[("size",(show . Map.size) x)], (fmap AnyData . Map.toList) x))

setCH :: CustomHandler
setCH = container1CH (\x -> ("Set",[("size",(show . Set.size) x)],(fmap AnyData . Set.toList) x))

intMapCH :: CustomHandler
intMapCH = container1CH (\x -> ("IntMap",[("size",(show . IntMap.size) x)],(fmap AnyData . IntMap.toList) x))

data AnyData = forall a. Data a => AnyData a

-- | Makes a 'CustomHandler' for container-like types. The given function should return:
--
-- * The string for the 'rowCV' cell
--
-- * A list of arbitrary (key,value) pairs to be displayed in the 'rowCustomInfo' cell (this should be things
-- like the size of the collection, not a list of elements)
--
-- * The list of elements
--
container0CH :: Typeable a =>
    (a -> (String,[(String,String)],[AnyData])) -> CustomHandler
container0CH convert = monoCH (\x -> containerCH_common (convert x) (typeOf x))

-- | Like 'container0CH', but for type constructors
container1CH :: Typeable1 f =>
    (forall a. (Data a) => f a -> (String,[(String,String)],[AnyData])) -> CustomHandler
container1CH convert = poly1CH (\x -> containerCH_common (convert x) (typeOf x))

-- | Like 'container1CH', but for binary type constructor
container2CH :: Typeable2 f => 
    (forall a b. (Data a, Data b) => f a b -> (String,[(String,String)],[AnyData])) -> CustomHandler
container2CH convert = poly2CH (\x -> containerCH_common (convert x) (typeOf x))


containerCH_common :: (String,[(String,String)],[AnyData]) -> TypeRep -> MCH (Maybe (StrictTree Row))
containerCH_common (name,attrs,asList) typeRep = do
            cs <- mapM (\(AnyData y) -> self y) asList
            let row = 
                 newRow [ txt $ concat ["{",name,"}"] ]
                        [ txt $ showAttrList attrs]
                        [ txt $ showType typeRep ] 

                        `addToAll` [ fgcolor' 0 0 0xA000 ] 

            (return . Just) (strictTree (row,cs))




showAttrList ::  [(String, String)] -> String
showAttrList = intercalate ", " . (fmap (\(k,v) -> k ++ "=" ++ v)) 


strCH ::  CustomHandler
strCH = monoPureCH' (\x -> {-# SCC "strCH/body" #-}
                  (strictTree (formatLit (show (x::String)) "" (typeOf x))
                        ))

-- | Strict ByteString
bsCH :: CustomHandler
bsCH = monoPureCH' (\x@(BSI.PS addr offset length_) -> 
                     strictTree (formatLit  
                        ((show . BS.unpack) x) 
                        (showAttrList [("addr",show addr) ,("offset",show offset) ,("length",show length_)])
                        (typeOf x))
                          )

-- | Lazy ByteString
blCH :: CustomHandler
blCH = container0CH (\x -> 
    let
        chunks = BL.toChunks x
    in
        ("Lazy ByteString", [("nchunks",(show . length) chunks)], fmap AnyData chunks)) 



funCH :: CustomHandler
funCH = poly2CH (\x -> 
                    let
                        _types = x undefined
                        color = fgcolor "darkorange"
                        row = newRow [txt "{function}", color ] 
                                     [] [txt $ showTypeOf x, color]
                    in
                         (return . Just . strictTree) row)


formatLit :: String -> String -> TypeRep -> Row
formatLit x info t = newRow [ txt $ x ] [ txt $ info ] [ txt $ showType t ] 
                     
                     `addToAll` [ fgcolor' 0 0x8000 0 ]




-- | Show a type with some prettification like replacing @[Char]@ with @String@.
showType :: TypeRep -> String
showType = replace "[Char]" "String" . show

showTypeOf ::  Typeable a => a -> String
showTypeOf = showType . typeOf

-- | A 'CustomHandler' that only works at a fixed type (and always works there), and generates a node with no children and with the default text in the /type/ column. 
-- The function should return the 'rowCV' in the first component and the 'rowCustomInfo' in the second. 
simpleCH :: Typeable a => (a -> MCH ([CellAttr],[CellAttr])) -> CustomHandler
simpleCH f = monoCH (\x -> do
    (y1,y2) <- f x 
    (return . Just) $ strictTree (newRow y1 y2 [txt $ showTypeOf x]))