{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# OPTIONS -Wall #-}
module DataTreeView.DataToTree where
import Data.Data
import Data.Monoid
import Data.String.Utils(replace)
import Control.Monad.IO.Control
import Control.Exception.Control
import Control.Applicative
import Control.Monad.Reader
import Prelude hiding(catch)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.IntMap as IntMap
import Data.List(intercalate)
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Internal as BSI
import qualified Data.ByteString.Lazy.Char8 as BL
import DataTreeView.Row
import DataTreeView.StrictTypes





-- | 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, MonadControlIO, Typeable) 

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

-- | 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 x = do 
            y <- {-# SCC "dataToTree/runCH" #-} runCH ch' x 
            z <- case y of
                      Just y' -> return y'
                      Nothing -> defaultHandler 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" ] ++)


defaultHandler :: forall e. Data e => e -> MCH (StrictTree Row)
defaultHandler x = do
    rec <- {-# SCC "defaultHandler/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 "defaultHandler/rec'" #-}
                  case safeConstrFields x of                                                    
                    Nothing -> rec                                                            
                    Just fieldNames ->                                                        
                     zipWith                                                                  
                      (\n b -> 
                         let nodeValue' = (nodeValue n) { rowFieldName = cellData [ txt $ b ] }
                         in strictTree (nodeValue', nodeChildren n))  
                        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

container0CH :: Typeable a =>
    (a -> (String,[(String,String)],[AnyData])) -> CustomHandler
container0CH convert = monoCH (\x -> containerCH_common (convert x) (typeOf x))

container1CH :: Typeable1 f =>
    (forall a. (Data a) => f a -> (String,[(String,String)],[AnyData])) -> CustomHandler
container1CH convert = poly1CH (\x -> containerCH_common (convert x) (typeOf x))

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]))