open Prelude

(* TODO: Decide on public interface. Probably should not include
   functions such as [unsafe_update_bal], but having [Balanced_batch]
   under [Account] also feels a bit awkward. *)

module Path = struct
  type t = string list [@@deriving compare, sexp]
end

(** Ensures that only accounts with valid type hierarchies can be constructed.
*)
module Kernel : sig
  type extra = { description : string }

  type t = private { type_ : Type.t; extra : extra; core : core }

  and core =
    (* Balance in some commodity *)
    | Leaf of Money.Commodity_id.t * Money.Diff.t
    | Subtree of t String.Map.t

  val make : Type.t -> extra -> core -> t option
end = struct
  type extra = { description : string }

  type t = { type_ : Type.t; extra : extra; core : core }

  and core =
    (* Balance in some commodity *)
    | Leaf of Money.Commodity_id.t * Money.Diff.t
    | Subtree of t String.Map.t

  let make type_ extra : core -> t option = function
    | Leaf (comm, bal) -> Some { type_; extra; core = Leaf (comm, bal) }
    | Subtree children ->
        if
          Map.for_all children ~f:(fun subacc ->
              Type.is_super subacc.type_ ~super:type_ ~strict:false)
        then Some { type_; extra; core = Subtree children }
        else None
end

type t = Kernel.t

type update_bal_error =
  | Empty_path
  | Unmatching_commodity_id of { in_account : Money.Commodity_id.t }
  | Not_a_leaf_account
  | Not_a_subtree_account
  | Not_found (* TODO: report at which level *)

(* We do not necessarily expect [aid] to be a valid path, as we
   always do for paths in the type hierarchy. The difference is that
   the type hierarchy is fixed, while the account hierarchy can
   change over the course of a year. *)
let rec unsafe_update_bal_aux (aid : Path.t) (dc : Money.Debit_credit.t)
    (by_amount : Money.Amount.t) (in_comm : Money.Commodity_id.t) (w : t) :
    (t, update_bal_error) result =
  match (aid, w.core) with
  | [], Leaf (comm, bal) ->
      if [%equal: Money.Commodity_id.t] in_comm comm then (* slay! *)
        let core' =
          Kernel.Leaf
            ( comm,
              Money.Diff.(bal + of_amount by_amount dc (Type.polarity w.type_))
            )
        in
        Ok (Option.value_exn (Kernel.make w.type_ w.extra core'))
      else (* bruh *)
        Error (Unmatching_commodity_id { in_account = comm })
  | [], Subtree _ -> Error Not_a_leaf_account
  | _ :: _, Leaf _ -> Error Not_a_subtree_account
  | aid0 :: aid', Subtree subaccs -> (
      match Map.find subaccs aid0 with
      | None -> Error Not_found
      | Some subacc ->
          let open Result.Let_syntax in
          let%map subacc' =
            (* TODO: when reporting at which level Not_found fails,
               we want to make sure that we extend the information
               in the error with the current aid0 (so we recover a
               full path to where the account is missing *)
            unsafe_update_bal_aux aid' dc by_amount in_comm subacc
          in
          let core' =
            Kernel.Subtree (Map.set subaccs ~key:aid0 ~data:subacc')
          in
          Option.value_exn (Kernel.make w.type_ w.extra core'))

(* Unfortunate but true, there has to be some kind of a root account
   with no type :) *)
type root = t String.Map.t

let unsafe_update_bal (aid : Path.t) (dc : Money.Debit_credit.t)
    (by_amount : Money.Amount.t) (in_comm : Money.Commodity_id.t) (w : root) :
    (root, update_bal_error) result =
  match aid with
  | [] -> Error Empty_path
  | aid0 :: aid' -> (
      match Map.find w aid0 with
      | None -> Error Not_found
      | Some subacc ->
          let open Result.Let_syntax in
          (* TODO: when reporting at which level Not_found fails,
               we want to make sure that we extend the information
               in the error with the current aid0 (so we recover a
               full path to where the account is missing *)
          let%map subacc' =
            unsafe_update_bal_aux aid' dc by_amount in_comm subacc
          in
          Map.set w ~key:aid0 ~data:subacc')

module Balanced_batch_acc_paths = Balanced_batch.Make (Path)

let apply_balanced_batch (b : Balanced_batch_acc_paths.t) (w : root) =
  Map.fold_result (Balanced_batch_acc_paths.entries b) ~init:w
    ~f:(fun ~key:aid ~(data : Balanced_batch_acc_paths.entry) w ->
      unsafe_update_bal aid data.dc data.amount data.commodity w)