Source

screentime-monitor / client.ml

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 () -> 
       force clear_string
       >>= fun cstring ->
       retry
         ~on_retry:(fun () -> printf "%sRetrying...\n" cstring)
         ~on_error:(fun err ->
           printf "%sFailed.  Will retry\n\n%s\n"
             cstring (Exn.sexp_of_t err |> Sexp.to_string_hum))
         (fun () -> try_with (fun () ->
            setup_conn user (fun user conn ->
              let stop = force on_term_signal in
              Clock.every' (sec 1.) ~stop (fun () ->
                let spinner = force spinner in
                Rpc.Rpc.dispatch Protocol.status conn user
                >>| function
                | Error err ->
                  printf "%s%s Unable to reach server\n\n%s\n"
                    cstring
                    (Error.sexp_of_t err |> Sexp.to_string_hum)
                    (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
            )))
       >>| function
       | Ok () -> ()
       | Error err ->
         printf "Failed.\n\n%s\n" (Exn.sexp_of_t err |> Sexp.to_string_hum)
    )
              
            

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
  ]