Source

screentime-monitor / client.ml

Full commit
open Core.Std
open Async.Std
module Shell = Core_extended.Std.Shell
open Client_common

let start =
  Command.async_basic
    ~summary:"Start the timer"
    Command.Spec.(
      shared_flags ()
      +> anon ("category" %: category_spec)
    )
    (fun username category () ->
       setup_conn username (fun username conn ->
         Rpc.Rpc.dispatch_exn Protocol.start conn (username,category)))

let stop =
  Command.async_basic
    ~summary:"Stop the timer"
    (shared_flags ())
    (fun username () -> setup_conn username (fun username conn ->
       Rpc.Rpc.dispatch_exn Protocol.stop conn username))

let span_string span =
  let { Time.Span.Parts.
          sign = _
      ; hr; min; sec
      ; ms = _ ; us = _
      } =
    Time.Span.to_parts span
  in
  sprintf "%02d:%02d:%02d" hr min sec
;;

let monitor =
  Command.async_basic
    ~summary:"Monitor your time spent"
    (shared_flags ())
    (fun user () -> setup_robust_conn ~reconnect_after:(sec 10.) user (fun user rconn ->
       force clear_string
       >>= fun cstring ->
       let stop = force on_term_signal in
       Clock.every' (sec 1.) ~stop (fun () ->
         let spinner = force spinner in
         match Robust_connection.get rconn with
         | None -> printf "%sWaiting for reconnection\n" cstring; return ()
         | Some conn ->
           try_with (fun () -> Rpc.Rpc.dispatch_exn Protocol.status conn user)
           >>| function
           | Error _ ->
             printf "%s%s Unable to reach server\n" cstring (spinner ())
           | Ok { Protocol.Status. state; elapsed } ->
             let total_elapsed =
               Map.data elapsed
               |> List.fold ~init:Time.Span.zero ~f:Time.Span.(+)
             in
             printf "%s%s %s %s\n"
               cstring
               (spinner ())
               (span_string total_elapsed)
               (match state with
                | `Active cat -> sprintf "{%s}" (Category.to_string cat)
                | `Not_running -> "(paused)"
               )
       );
       stop
     ))

module Ascii_table = Textutils.Ascii_table

let report =
  Command.async_basic
    ~summary:"Print a report of time spent"
    Command.Spec.(
      shared_flags ()
      +> days_flag ()
    )
    (fun user days () -> setup_conn user (fun user conn ->
       Rpc.Rpc.dispatch_exn Protocol.report conn user
       >>= fun report ->
       let report =
         Map.to_alist report
         |> (fun l -> List.drop l (List.length l - days))
       in
       let categories =
         List.map ~f:snd report
         |> List.map ~f:Map.keys
         |> List.concat
         |> List.dedup
       in
       let open Ascii_table in
       let date_column =
         Column.create "date"
           (fun (date,_) -> Date.to_string date)
       in
       let cat_columns =
         List.map categories ~f:(fun cat ->
           Column.create
             (Category.to_string cat)
             (fun (_date,by_cat) ->
                match Map.find by_cat cat with
                | None -> "-"
                | Some span -> span_string span
             ))
       in
       let total_column =
         Column.create "total"
           (fun (_date,by_cat) ->
              Map.data by_cat
              |> List.fold ~init:Time.Span.zero ~f:Time.Span.(+)
              |> span_string
           )
       in
       printf "%s\n"
         (to_string
            ~display:Display.column_titles
            ([date_column]
             @ (if List.length categories > 1 then cat_columns else [])
             @ [total_column]
            )
            report);
       return ()
     ))

let full_report =
  Command.async_basic
    ~summary:"Print a detailed report of all the screentime sessions"
    Command.Spec.(
      shared_flags ()
      +> days_flag ()
    )
    (fun user days () -> setup_conn user (fun user conn ->
       Rpc.Rpc.dispatch_exn Protocol.full_report conn user
       >>= fun report ->
       let open Ascii_table in
       let ofday_string t =
         Time.to_local_ofday t |> Time.Ofday.to_sec_string
       in
       let columns =
         [ Column.create "category" (fun (c,_,_) -> Category.to_string c)
         ; Column.create "start"    (fun (_,t,_) -> ofday_string t)
         ; Column.create "stop"     (fun (_,_,t) -> ofday_string t)
         ; Column.create "min"     (fun (_,start,stop) ->
           Time.diff stop start
           |> Time.Span.to_min
           |> Float.to_int
           |> Int.to_string)
         ]
       in
       Map.to_alist report
       |> (fun l -> List.drop l (List.length l - days))
       |> List.iter ~f:(fun (date,sessions) ->
         printf "*** Date: %s ***\n\n" (Date.to_string date);
         printf "%s\n" (to_string
                          ~display:Display.column_titles
                          columns
                          (List.rev sessions)));
       Deferred.unit
     ))

let commands =
  [ "start"       , start
  ; "stop"        , stop
  ; "monitor"     , monitor
  ; "report"      , report
  ; "full-report" , full_report
  ; "rules"       , Rule_client.command
  ]