Fix JSON dump field names (fixes #45)
[darcs-mirror-arbtt.git] / src / Categorize.hs
1 {-# LANGUAGE Rank2Types, CPP #-}
2 module Categorize where
3
4 import Data
5
6 import qualified Text.Regex.PCRE.Light.Text as RE
7 import qualified Data.Map as M
8 import qualified Data.MyText as T
9 import Data.MyText (Text)
10 import Control.Monad
11 import Control.Monad.Instances
12 import Control.Monad.Trans.Reader
13 import Control.Monad.Trans.Class
14 import Data.Functor.Identity
15
16
17 import Text.Parsec
18 import Text.Parsec.Char
19 import Text.Parsec.Token
20 import Text.Parsec.Combinator
21 import Text.Parsec.Language
22 import Text.Parsec.ExprFail
23 import System.IO
24 import System.Exit
25 import Control.Applicative ((<*>),(<$>))
26 import Control.DeepSeq
27 import Data.List
28 import Data.Maybe
29 import Data.Char
30 import Data.Time.Clock
31 import Data.Time.LocalTime
32 import Data.Time.Calendar (toGregorian, fromGregorian)
33 import Data.Time.Calendar.WeekDate (toWeekDate)
34 import Data.Time.Format (formatTime)
35 #if MIN_VERSION_time(1,5,0)
36 import Data.Time.Format(defaultTimeLocale, iso8601DateFormat)
37 #else
38 import System.Locale (defaultTimeLocale, iso8601DateFormat)
39 #endif
40 import Debug.Trace
41 import Control.Arrow (second)
42 import Text.Printf
43
44 type Categorizer = TimeLog CaptureData -> TimeLog (Ctx, ActivityData)
45 type Rule = Ctx -> ActivityData
46
47 type Parser = ParsecT String () (ReaderT TimeZone Identity)
48
49
50 data Ctx = Ctx
51         { cNow :: TimeLogEntry CaptureData
52         , cCurrentWindow :: Maybe (Bool, Text, Text)
53         , cWindowInScope :: Maybe (Bool, Text, Text)
54         , cSubsts :: [Text]
55         , cCurrentTime :: UTCTime
56         , cTimeZone :: TimeZone
57         }
58   deriving (Show)
59
60 instance NFData Ctx where
61     rnf (Ctx a b c d e f) = a `deepseq` b `deepseq` c `deepseq` e `deepseq` e `deepseq` f `deepseq` ()
62
63 type Cond = CtxFun [Text]
64
65 type CtxFun a = Ctx -> Maybe a
66
67 data CondPrim
68         = CondString (CtxFun Text)
69         | CondRegex (CtxFun RE.Regex)
70         | CondInteger (CtxFun Integer)
71         | CondTime (CtxFun NominalDiffTime)
72         | CondDate (CtxFun UTCTime)
73         | CondCond (CtxFun [Text])
74         | CondStringList (CtxFun [Text])
75         | CondRegexList (CtxFun [RE.Regex])
76
77 newtype Cmp = Cmp (forall a. Ord a => a -> a -> Bool)
78
79 readCategorizer :: FilePath -> IO Categorizer
80 readCategorizer filename = do
81         h <- openFile filename ReadMode
82         hSetEncoding h utf8
83         content <- hGetContents h
84         time <- getCurrentTime
85         tz <- getCurrentTimeZone
86         case flip runReader tz $
87             runParserT (between (return ()) eof parseRules) () filename content of
88           Left err -> do
89                 putStrLn "Parser error:"
90                 print err
91                 exitFailure
92           Right cat -> return $
93                 (map (fmap (mkSecond (postpare . cat))) . prepare time tz)
94
95 applyCond :: String -> TimeZone -> TimeLogEntry (Ctx, ActivityData) -> Bool
96 applyCond s tz = 
97         case flip runReader tz $ runParserT (do {c <- parseCond; eof ; return c}) () "commad line parameter" s of
98           Left err -> error (show err)
99           Right c  -> isJust . c . fst . tlData
100
101 prepare :: UTCTime -> TimeZone -> TimeLog CaptureData -> TimeLog Ctx
102 prepare time tz = map go
103   where go now  = now {tlData = Ctx now (findActive (cWindows (tlData now))) Nothing [] time tz }
104
105 -- | Here, we filter out tags appearing twice, and make sure that only one of
106 --   each category survives
107 postpare :: ActivityData -> ActivityData
108 postpare = nubBy go
109   where go (Activity (Just c1) _) (Activity (Just c2) _) = c1 == c2
110         go a1                     a2                     = a1 == a2
111
112 lang :: GenTokenParser String () (ReaderT TimeZone Identity)
113 lang = makeTokenParser $ LanguageDef
114                 { commentStart   = "{-"
115                 , commentEnd     = "-}"
116                 , commentLine    = "--"
117                 , nestedComments = True
118                 , identStart     = letter
119                 , identLetter    = alphaNum <|> oneOf "_'"
120                 , opStart        = oneOf ":!#$%&*+./<=>?@\\^|-~"
121                 , opLetter       = oneOf ":!#$%&*+./<=>?@\\^|-~"
122                 , reservedOpNames= []
123                 , reservedNames  = []
124                 , caseSensitive  = True
125                 }
126
127 parseRules :: Parser Rule
128 parseRules = do 
129         whiteSpace lang
130         a <- option id (reserved lang "aliases" >> parens lang parseAliasSpecs)
131         rb <- parseRulesBody
132         return (a . rb)
133
134 parseAliasSpecs :: Parser (ActivityData -> ActivityData)
135 parseAliasSpecs = do as <- sepEndBy1 parseAliasSpec (comma lang)
136                      return $ \ad -> foldr doAlias ad as
137
138 doAlias :: (Text, Text) -> ActivityData -> ActivityData
139 doAlias (s1,s2) = map go
140   where go (Activity cat tag) = Activity (if cat == Just s1 then Just s2 else cat)
141                                          (if tag == s1 then s2 else tag)
142
143 parseAliasSpec :: Parser (Text, Text)
144 parseAliasSpec = do s1 <- T.pack <$> stringLiteral lang
145                     reservedOp lang "->"
146                     s2 <- T.pack <$> stringLiteral lang
147                     return (s1,s2)
148
149 parseRulesBody :: Parser Rule
150 parseRulesBody = do 
151         x <- parseRule
152         choice [ do comma lang
153                     xs <- parseRule `sepEndBy1` comma lang
154                     return (matchAny (x:xs))
155                , do semi lang
156                     xs <- parseRule `sepEndBy1` semi lang
157                     return (matchFirst (x:xs))
158                ,    return x
159                ]
160
161 parseRule :: Parser Rule
162 parseRule = choice
163         [    braces lang parseRules
164         , do cond <- parseCond
165              reservedOp lang "==>"
166              rule <- parseRule
167              return (ifThenElse cond rule matchNone)
168         , do reserved lang "if"
169              cond <- parseCond
170              reserved lang "then"
171              rule1 <- parseRule
172              reserved lang "else"
173              rule2 <- parseRule
174              return (ifThenElse cond rule1 rule2)
175         , do reserved lang "tag"
176              parseSetTag
177         ]
178
179 parseCond :: Parser Cond
180 parseCond = do cp <- parseCondExpr
181                case cp of
182                 CondCond c -> return c
183                 cp         -> fail $ printf "Expression of type %s" (cpType cp)
184
185 parseCondExpr :: Parser CondPrim
186 parseCondExpr = buildExpressionParser [
187                 [ Prefix (reservedOp lang "!" >> return checkNot) ],
188                 [ Prefix (reserved lang "day of week" >> return evalDayOfWeek)
189                 , Prefix (reserved lang "day of month" >> return evalDayOfMonth)
190                 , Prefix (reserved lang "month" >> return evalMonth)
191                 , Prefix (reserved lang "year" >> return evalYear)
192                 , Prefix (reserved lang "format" >> return formatDate) ],
193                 [ Infix (reservedOp lang "=~" >> return checkRegex) AssocNone 
194                 , Infix (checkCmp <$> parseCmp) AssocNone
195                 ],
196                 [ Prefix (reserved lang "current window" >> return checkCurrentwindow)
197                 , Prefix (reserved lang "any window" >> return checkAnyWindow)
198                 ],
199                 [ Infix (reservedOp lang "&&" >> return checkAnd) AssocRight ],
200                 [ Infix (reservedOp lang "||" >> return checkOr) AssocRight ]
201             ] parseCondPrim
202
203 cpType :: CondPrim -> String
204 cpType (CondString _) = "String"
205 cpType (CondRegex _) = "Regex"
206 cpType (CondInteger _) = "Integer"
207 cpType (CondTime _) = "Time"
208 cpType (CondDate _) = "Date"
209 cpType (CondCond _) = "Condition"
210 cpType (CondStringList _) = "List of Strings"
211 cpType (CondRegexList _) = "List of regular expressions"
212
213 checkRegex :: CondPrim -> CondPrim -> Erring CondPrim
214 checkRegex (CondString getStr) (CondRegex getRegex) = Right $ CondCond $ \ctx -> do
215         str <- getStr ctx
216         regex <- getRegex ctx
217         tail <$> RE.match regex str [RE.exec_no_utf8_check]
218 checkRegex (CondString getStr) (CondRegexList getRegexList) = Right $ CondCond $ \ctx -> do
219         str <- getStr ctx
220         regexes <- getRegexList ctx
221         tail <$> msum (map (\regex -> RE.match regex str [RE.exec_no_utf8_check]) regexes)
222 checkRegex cp1 cp2 = Left $
223         printf "Cannot apply =~ to an expression of type %s and type %s"
224                (cpType cp1) (cpType cp2)
225
226 checkAnd :: CondPrim-> CondPrim -> Erring CondPrim
227 checkAnd (CondCond c1) (CondCond c2) = Right $ CondCond $ do
228         res1 <- c1
229         res2 <- c2
230         return $ res1 >> res2
231 checkAnd cp1 cp2 = Left $
232         printf "Cannot apply && to an expression of type %s and type %s"
233                (cpType cp1) (cpType cp2)
234
235 checkOr :: CondPrim-> CondPrim -> Erring CondPrim
236 checkOr (CondCond c1) (CondCond c2) = Right $ CondCond $ do
237         res1 <- c1
238         res2 <- c2
239         return $ res1 `mplus` res2
240 checkOr cp1 cp2 = Left $
241         printf "Cannot apply && to an expression of type %s and type %s"
242                (cpType cp1) (cpType cp2)
243
244 checkNot :: CondPrim -> Erring CondPrim
245 checkNot (CondCond getCnd) = Right $ CondCond $ do
246         liftM (maybe (Just []) (const Nothing)) getCnd
247 checkNot cp = Left $
248         printf "Cannot apply ! to an expression of type %s"
249                (cpType cp)
250
251 checkCmp :: Cmp -> CondPrim -> CondPrim -> Erring CondPrim
252 checkCmp (Cmp (?)) (CondInteger getN1) (CondInteger getN2) = Right $ CondCond $ \ctx -> do
253         n1 <- getN1 ctx
254         n2 <- getN2 ctx
255         guard (n1 ? n2)
256         return []
257 checkCmp (Cmp (?)) (CondTime getT1) (CondTime getT2) = Right $ CondCond $ \ctx -> do
258         t1 <- getT1 ctx
259         t2 <- getT2 ctx
260         guard (t1 ? t2)
261         return []
262 checkCmp (Cmp (?)) (CondDate getT1) (CondDate getT2) = Right $ CondCond $ \ctx -> do
263         t1 <- getT1 ctx
264         t2 <- getT2 ctx
265         guard (t1 ? t2)
266         return []
267 checkCmp (Cmp (?)) (CondString getS1) (CondString getS2) = Right $ CondCond $ \ctx -> do
268         s1 <- getS1 ctx
269         s2 <- getS2 ctx
270         guard (s1 ? s2)
271         return []
272 checkCmp (Cmp (?)) (CondString getS1) (CondStringList getS2) = Right $ CondCond $ \ctx -> do
273         s1 <- getS1 ctx
274         sl <- getS2 ctx
275         guard (any (s1 ?) sl)
276         return []
277 checkCmp _ cp1 cp2 = Left $
278         printf "Cannot compare expressions of type %s and type %s"
279                (cpType cp1) (cpType cp2)
280
281 checkCurrentwindow :: CondPrim -> Erring CondPrim
282 checkCurrentwindow (CondCond cond) = Right $ CondCond $ \ctx -> 
283         cond (ctx { cWindowInScope = cCurrentWindow ctx })
284 checkCurrentwindow cp = Left $
285         printf "Cannot apply current window to an expression of type %s"
286                (cpType cp)
287
288 checkAnyWindow :: CondPrim -> Erring CondPrim
289 checkAnyWindow (CondCond cond) = Right $ CondCond $ \ctx ->
290         msum $ map (\w -> cond (ctx { cWindowInScope = Just w }))
291                                      (cWindows (tlData (cNow ctx)))
292 checkAnyWindow cp = Left $
293         printf "Cannot apply current window to an expression of type %s"
294                (cpType cp)
295
296 fst3 (a,_,_) = a
297 snd3 (_,b,_) = b
298 trd3 (_,_,c) = c
299
300 -- Day of week is an integer in [1..7].
301 evalDayOfWeek :: CondPrim -> Erring CondPrim
302 evalDayOfWeek (CondDate df) = Right $ CondInteger $ \ctx ->
303   let tz = cTimeZone ctx in
304   (toInteger . trd3 . toWeekDate . localDay . utcToLocalTime tz) `liftM` df ctx
305 evalDayOfWeek cp = Left $ printf
306   "Cannot apply day of week to an expression of type %s, only to $date."
307   (cpType cp)
308
309 -- Day of month is an integer in [1..31].
310 evalDayOfMonth :: CondPrim -> Erring CondPrim
311 evalDayOfMonth (CondDate df) = Right $ CondInteger $ \ctx ->
312   let tz = cTimeZone ctx in
313   (toInteger . trd3 . toGregorian . localDay . utcToLocalTime tz) `liftM` df ctx
314 evalDayOfMonth cp = Left $ printf
315   "Cannot apply day of month to an expression of type %s, only to $date."
316   (cpType cp)
317
318 -- Month is an integer in [1..12].
319 evalMonth :: CondPrim -> Erring CondPrim
320 evalMonth (CondDate df) = Right $ CondInteger $ \ctx ->
321   let tz = cTimeZone ctx in
322   (toInteger . snd3 . toGregorian . localDay . utcToLocalTime tz) `liftM` df ctx
323 evalMonth cp = Left $ printf
324   "Cannot apply month to an expression of type %s, only to $date."
325   (cpType cp)
326
327 evalYear :: CondPrim -> Erring CondPrim
328 evalYear (CondDate df) = Right $ CondInteger $ \ctx ->
329   let tz = cTimeZone ctx in
330   (fst3 . toGregorian . localDay . utcToLocalTime tz) `liftM` df ctx
331 evalYear cp = Left $ printf
332   "Cannot apply year to an expression of type %s, only to $date."
333   (cpType cp)
334
335 -- format date according to ISO 8601 (YYYY-MM-DD)
336 formatDate :: CondPrim -> Erring CondPrim
337 formatDate (CondDate df) = Right $ CondString $ \ctx ->
338   let tz = cTimeZone ctx
339       local = utcToLocalTime tz `liftM` df ctx
340    in T.pack . formatTime defaultTimeLocale (iso8601DateFormat Nothing) <$> local
341 formatDate cp = Left $ printf
342   "Cannot format an expression of type %s, only $date." (cpType cp)
343
344 parseCmp :: Parser Cmp
345 parseCmp = choice $ map (\(s,o) -> reservedOp lang s >> return o)
346                         [(">=",Cmp (>=)),
347                          (">", Cmp (>)),
348                          ("==",Cmp (==)),
349                          ("=", Cmp (==)),
350                          ("!=",Cmp (/=)),
351                          ("<", Cmp (<)),
352                          ("<=",Cmp (<=))]
353
354 parseCondPrim :: Parser CondPrim
355 parseCondPrim = choice
356         [ parens lang parseCondExpr
357         , brackets lang (choice [
358             (do list <- commaSep1 lang (stringLiteral lang)
359                 return $ CondStringList (const (Just (map T.pack list)))
360             ) <?> "list of strings",
361             (do list <- commaSep1 lang parseRegex
362                 return $ CondRegexList (const (Just list))
363             ) <?> "list of regular expressions"
364             ])
365         , char '$' >> choice 
366              [ do backref <- read <$> many1 digit
367                   return $ CondString (getBackref backref)
368              , do varname <- identifier lang 
369                   choice 
370                       [ do guard $ varname == "title"
371                            return $ CondString (getVar "title")
372                       , do guard $ varname == "program"
373                            return $ CondString (getVar "program")
374                       , do guard $ varname == "active"
375                            return $ CondCond checkActive
376                       , do guard $ varname == "idle"
377                            return $ CondInteger (getNumVar "idle")
378                       , do guard $ varname == "time"
379                            return $ CondTime (getTimeVar "time")
380                       , do guard $ varname == "sampleage"
381                            return $ CondTime (getTimeVar "sampleage")
382                       , do guard $ varname == "date"
383                            return $ CondDate (getDateVar "date")
384                       , do guard $ varname == "now"
385                            return $ CondDate (getDateVar "now")
386                       , do guard $ varname == "desktop"
387                            return $ CondString (getVar "desktop")
388                      ]
389               ] <?> "variable"
390         , do regex <- parseRegex <?> "regular expression"
391              return $ CondRegex (const (Just regex))
392         , do str <- T.pack <$> stringLiteral lang <?> "string"
393              return $ CondString (const (Just str))
394         , try $ do time <- parseTime <?> "time" -- backtrack here, it might have been a number
395                    return $ CondTime (const (Just time))
396         , try $ do date <- parseDate <?> "date" -- backtrack here, it might have been a number
397                    return $ CondDate (const (Just date))
398         , do num <- natural lang <?> "number"
399              return $ CondInteger (const (Just num))
400         ]
401 {-
402                      choice
403                         [ do reservedOp lang "=~"
404                              regex <- parseRegex
405                              return $ checkRegex varname (RE.compile regex [])
406                         , do reservedOp lang "==" <|> reservedOp lang "="
407                              str <- stringLiteral lang
408                              return $ checkEq varname str
409                         , do reservedOp lang "/=" <|> reservedOp lang "!="
410                              str <- stringLiteral lang
411                              return $ checkNot (checkEq varname str)
412                         ]
413                 , do guard $ varname == "idle"
414                      op <- parseCmp
415                      num <- natural lang
416                      return $ checkNumCmp op varname num
417                 , do guard $ varname `elem` ["time","sampleage"]
418                      op <- parseCmp 
419                      time <- parseTime
420                      return $ checkTimeCmp op varname time
421                 , do guard $ varname == "active"
422                      return checkActive
423                 ]
424         , do reserved lang "current window"
425              cond <- parseCond
426              return $ checkCurrentwindow cond
427         , do reserved lang "any window"
428              cond <- parseCond
429              return $ checkAnyWindow cond
430         ]
431 -}
432
433 parseRegex :: Parser RE.Regex
434 parseRegex = fmap (flip RE.compile [] . T.pack) $ lexeme lang $ choice
435         [ between (char '/') (char '/') (many1 (noneOf "/"))
436         , do char 'm'
437              c <- anyChar
438              str <- many1 (noneOf [c])
439              char c
440              return str
441         ]
442              
443 -- | Parses a day-of-time specification (hh:mm)
444 parseTime :: Parser NominalDiffTime
445 parseTime = fmap fromIntegral $ lexeme lang $ do
446                hour <- read <$> many1 digit
447                char ':'
448                minute <- read <$> count 2 digit
449                return $ (hour * 60 + minute) * 60
450
451 parseDate :: Parser UTCTime
452 parseDate = lexeme lang $ do
453     tz <- lift ask
454     year <- read <$> count 4 digit
455     char '-'
456     month <- read <$> count 2 digit
457     char '-'
458     day <- read <$> count 2 digit
459     time <- option 0 parseTime
460     let date = LocalTime (fromGregorian year month day) (TimeOfDay 0 0 0)
461     return $ addUTCTime time $ localTimeToUTC tz date
462
463
464 parseSetTag :: Parser Rule
465 parseSetTag = lexeme lang $ do
466                  firstPart <- parseTagPart 
467                  choice [ do char ':'
468                              secondPart <- parseTagPart
469                              return $ do cat <- firstPart
470                                          tag <- secondPart
471                                          return $ maybeToList $ do
472                                             cat <- cat
473                                             tag <- tag
474                                             return $ Activity (Just cat) tag
475                         ,    return $ do tag <- firstPart
476                                          return $ maybeToList $ do
477                                             tag <- tag
478                                             return $ Activity Nothing tag
479                         ]
480
481 replaceForbidden :: Maybe Text -> Maybe Text
482 replaceForbidden = liftM $ T.map go
483   where
484     go c | isAlphaNum c  = c
485          | c `elem` "-_" = c
486          | otherwise     = '_'
487
488 parseTagPart :: Parser (Ctx -> Maybe Text)
489 parseTagPart = do parts <- many1 (choice
490                         [ do char '$'
491                              (replaceForbidden . ) <$> choice
492                                [ do num <- read <$> many1 digit
493                                     return $ getBackref num
494                                , do varname <- many1 (letter <|> oneOf ".")
495                                     return $ getVar varname
496                                ] <?> "variable"
497                         , do s <- many1 (alphaNum <|> oneOf "-_")
498                              return $ const (Just (T.pack s))
499                         ])
500                   return $ (fmap T.concat . sequence) <$> sequence parts
501
502 ifThenElse :: Cond -> Rule -> Rule -> Rule
503 ifThenElse cond r1 r2 = do res <- cond
504                            case res of 
505                             Just substs -> r1 . setSubsts substs
506                             Nothing -> r2
507   where setSubsts :: [Text] -> Ctx -> Ctx
508         setSubsts substs ctx = ctx { cSubsts = substs }
509         
510
511 matchAny :: [Rule] -> Rule
512 matchAny rules = concat <$> sequence rules
513 matchFirst :: [Rule] -> Rule
514 matchFirst rules = takeFirst <$> sequence rules
515   where takeFirst [] = []
516         takeFirst ([]:xs) = takeFirst xs
517         takeFirst (x:xs) = x
518
519
520 getBackref :: Integer -> CtxFun Text
521 getBackref n ctx = listToMaybe (drop (fromIntegral n-1) (cSubsts ctx))
522
523 getVar :: String -> CtxFun Text
524 getVar v ctx | "current" `isPrefixOf` v = do
525                 let var = drop (length "current.") v
526                 win <- cCurrentWindow ctx
527                 getVar var (ctx { cWindowInScope = Just win })
528 getVar "title"   ctx = do
529                 (_,t,_) <- cWindowInScope ctx
530                 return t
531 getVar "program" ctx = do
532                 (_,_,p) <- cWindowInScope ctx
533                 return p
534 getVar "desktop" ctx = do
535                 return $ cDesktop (tlData (cNow ctx))
536 getVar v ctx = error $ "Unknown variable " ++ v
537
538 getNumVar :: String -> CtxFun Integer
539 getNumVar "idle" ctx = Just $ cLastActivity (tlData (cNow ctx)) `div` 1000
540
541 getTimeVar :: String -> CtxFun NominalDiffTime
542 getTimeVar "time" ctx = Just $
543    let utc = tlTime . cNow $ ctx
544        tz = cTimeZone ctx
545        local = utcToLocalTime tz utc
546        midnightUTC = localTimeToUTC tz $ local { localTimeOfDay = midnight }
547     in utc `diffUTCTime` midnightUTC
548 getTimeVar "sampleage" ctx = Just $ cCurrentTime ctx `diffUTCTime` tlTime (cNow ctx)
549
550 getDateVar :: String -> CtxFun UTCTime
551 getDateVar "date" = Just . tlTime . cNow
552 getDateVar "now" = Just . cCurrentTime
553
554 findActive :: [(Bool, t, t1)] -> Maybe (Bool, t, t1)
555 findActive = find (\(a,_,_) -> a)                                 
556
557 checkActive :: Cond
558 checkActive ctx = do (a,_,_) <- cWindowInScope ctx
559                      guard a
560                      return []
561
562 matchNone :: Rule
563 matchNone = const []
564
565 justIf :: a -> Bool -> Maybe a
566 justIf x True = Just x
567 justIf x False = Nothing
568
569 mkSecond :: (a -> b) -> a -> (a, b)
570 mkSecond f a = (a, f a)