open Core
open Stdio
open Base

type mood = Val.t
type t = (Date.t, mood) Ordered_map.t

let now () = Date.today ~zone:Time_float.Zone.utc
let mood_of_day t day = Ordered_map.find t ~key:day
let create () = Ordered_map.create ~cmp:Core.Date.compare
let record_day t ~day ~mood = Ordered_map.add t ~key:day ~data:mood

let record ?day t ~mood =
  record_day t ~day:(Option.value day ~default:(now ())) ~mood

let today t =
  let now = now () in
  (now, mood_of_day t now)

let this_week t =
  let now = now () in
  let day_of_week = Day_of_week.to_int @@ Date.day_of_week now in
  let monday = Date.add_days now (-day_of_week) in
  let week = List.init 7 ~f:(fun i -> Date.add_days monday i) in
  List.map week ~f:(fun d -> (d, mood_of_day t d))

let this_month t =
  let now = now () in
  let y = Date.year now in
  let m = Date.month now in
  let first_day = Date.create_exn ~y ~m ~d:1 in
  let days_in_month = Date.days_in_month ~year:y ~month:m in
  let month = List.init days_in_month ~f:(fun i -> Date.add_days first_day i) in
  List.map month ~f:(fun d -> (d, mood_of_day t d))

let this_year t =
  let now = now () in
  let y = Date.year now in
  let first_day = Date.create_exn ~y ~m:Month.Jan ~d:1 in
  let days_in_year = if Date.is_leap_year ~year:y then 366 else 365 in
  let year = List.init days_in_year ~f:(fun i -> Date.add_days first_day i) in
  List.map year ~f:(fun d -> (d, mood_of_day t d))

let all t =
  let acc = ref [] in
  Ordered_map.iter t ~f:(fun ~key ~data -> acc := (key, data) :: !acc);
  List.rev !acc

let entry_of_string l =
  match String.split_on_chars ~on:[ ',' ] l with
  | [ date; v ] -> Some (Date.of_string date, Val.of_string v)
  | _ -> None

let entry_to_string (day, mood) =
  String.concat ~sep:"," [ Date.to_string day; Val.to_string mood ]

let of_list ls =
  List.fold ~init:(create ())
    ~f:(fun acc (day, mood) -> record_day acc ~day ~mood)
    ls

let to_list = all

let load location =
  try
    In_channel.with_file location ~f:(fun c ->
        List.filter_map (In_channel.input_lines c) ~f:entry_of_string)
    |> of_list
  with _ -> create ()

let store t ~location =
  Out_channel.write_lines location (List.map ~f:entry_to_string (all t))

let pp t = String.concat ~sep:"\n" @@ List.map ~f:entry_to_string (all t)

let random () =
  Random.init 12334254;
  let now = now () in
  let y = Date.year now in
  let first_day = Date.create_exn ~y ~m:Month.Jan ~d:1 in
  let days_in_year = if Date.is_leap_year ~year:y then 366 else 365 in
  let year = List.init days_in_year ~f:(fun i -> Date.add_days first_day i) in
  List.map year ~f:(fun d -> (d, Val.of_int (Random.int 5))) |> of_list

let render t =
  let vals = this_year t in
  let now = now () in
  let y = Date.year now in
  let start_date = Date.create_exn ~y ~m:Month.Jan ~d:1 in

  let start_dow = Date.day_of_week start_date |> Day_of_week.to_int in
  let pad_tracker n =
    List.init n ~f:(fun _ -> (Date.of_string "2000-01-01", None))
  in
  let start_padding = pad_tracker start_dow in
  let total_len = List.length start_padding + List.length vals in
  let remainder = total_len % 7 in
  let end_padding = pad_tracker (7 - remainder) in
  let padded = start_padding @ vals @ end_padding in
  let weeks = List.chunks_of padded ~length:7 in
  let rows = List.transpose_exn weeks in
  let day_labels = [ "Sun"; ""; "Tue"; ""; "Thu"; ""; "Sat" ] in
  List.map2_exn day_labels rows ~f:(fun label row ->
      let blocks =
        List.map row ~f:(fun (_, mood) ->
            Option.value_map mood
              ~default:("\027[2m" ^ "·" ^ "\027[0m")
              ~f:Val.render)
        |> String.concat
      in
      sprintf "%-3s %s" label blocks)
  |> String.concat ~sep:"\n"