open Prelude

type tx_type =
  | Interest_tx
  | Online_banking_tx
  | Recurrent_direct_tx
  | Payment_terminal_tx
  | Cash_payment_tx
  | Atm_tx
  | Auto_save_rounding_tx
  | Batch_tx
  | Direct_debit_tx
  | Periodic_tx

type iban_tag = Account_tag | Counterparty_iban_tag [@@deriving compare, sexp]

type unit_tag = Filed_tag | Google_pay_tag | Auto_round_savings_tag
[@@deriving compare, sexp]

type string_tag =
  | Desc_tag
  | User_tag
  | Counterparty_name_tag
  | Reference_tag
  | Mandate_id_tag
  | Creditor_id_tag
  | Other_party_tag
  | Transaction_tag
  | Terminal_tag
  | Card_seq_no_tag
  | Savings_account_tag
[@@deriving compare, sexp]

module Label = struct
  type 'a t =
    | Iban_label : iban_tag -> Iban.t t
    | String_label : string_tag -> string t
    | Timestamp_label : Time_ns.t t
    | Unit_label : unit_tag -> unit t

  let int_to_cmp x : ('a, 'a) Dmap.cmp =
    if x < 0 then Lt else if x > 0 then Gt else Eq

  let compare (type a1 a2) (v1 : a1 t) (v2 : a2 t) : (a1, a2) Dmap.cmp =
    match (v1, v2) with
    | Iban_label t1, Iban_label t2 -> int_to_cmp @@ [%compare: iban_tag] t1 t2
    | String_label t1, String_label t2 ->
        int_to_cmp @@ [%compare: string_tag] t1 t2
    | Timestamp_label, Timestamp_label -> Eq
    | Unit_label t1, Unit_label t2 -> int_to_cmp @@ [%compare: unit_tag] t1 t2
    | Iban_label _, _ -> Lt
    | String_label _, Iban_label _ -> Gt
    | String_label _, _ -> Lt
    | Timestamp_label, Unit_label _ -> Lt
    | Timestamp_label, _ -> Gt
    | Unit_label _, _ -> Gt
end

module Labels = struct
  include Dmap.Make (Label)

  let sexp_of_binding = function
    | Binding (Iban_label tag, iban) ->
        Sexp.List
          [
            Sexp.Atom "iban"; [%sexp_of: iban_tag] tag; [%sexp_of: Iban.t] iban;
          ]
    | Binding (String_label tag, s) ->
        Sexp.List
          [ Sexp.Atom "string"; [%sexp_of: string_tag] tag; Sexp.Atom s ]
    | Binding (Timestamp_label, ts) ->
        Sexp.List [ Sexp.Atom "timestamp"; [%sexp_of: Time_ns_unix.t] ts ]
    | Binding (Unit_label tag, ()) ->
        Sexp.List [ Sexp.Atom "unit"; [%sexp_of: unit_tag] tag ]

  let binding_of_sexp sexp =
    match sexp with
    | Sexp.List [ Sexp.Atom "iban"; tag_sexp; iban_sexp ] ->
        Binding
          ( Iban_label ([%of_sexp: iban_tag] tag_sexp),
            [%of_sexp: Iban.t] iban_sexp )
    | Sexp.List [ Sexp.Atom "string"; tag_sexp; Sexp.Atom s ] ->
        Binding (String_label ([%of_sexp: string_tag] tag_sexp), s)
    | Sexp.List [ Sexp.Atom "timestamp"; ts_sexp ] ->
        Binding (Timestamp_label, [%of_sexp: Time_ns_unix.t] ts_sexp)
    | Sexp.List [ Sexp.Atom "unit"; tag_sexp ] ->
        Binding (Unit_label ([%of_sexp: unit_tag] tag_sexp), ())
    | _ -> of_sexp_error "Labels.binding_of_sexp: invalid binding" sexp

  let sexp_of_t m = Sexp.List (bindings m |> List.map ~f:sexp_of_binding)

  let t_of_sexp sexp =
    match sexp with
    | Sexp.List labels ->
        Sequence.(of_list labels >>| binding_of_sexp |> to_seq) |> of_seq
    | Sexp.Atom _ -> of_sexp_error "Labels.t_of_sexp: list needed" sexp
end

module Debit_credit = struct
  type t = Debit | Credit [@@deriving string, sexp_of]

  (*  let opposite = function Debit -> Credit | Credit -> Debit *)
end

module Money = struct
  module Amount : sig
    type t

    val equal : t -> t -> bool
    val compare : t -> t -> int
    val of_bigint : Bigint.t -> t option
    val to_bigint : t -> Bigint.t
    val ( + ) : t -> t -> t
    val ( = ) : t -> t -> bool
    val sexp_of_t : t -> Sexp.t
    val zero : t
  end = struct
    type t = Bigint.t [@@deriving sexp_of]

    let equal = Bigint.equal
    let compare = Bigint.compare
    let of_bigint x = if Bigint.(zero <= x) then Some x else None
    let to_bigint x = x
    let ( + ) x y = Bigint.(x + y)
    let ( = ) = equal
    let zero = Bigint.zero
  end

  module Diff : sig
    type t

    val equal : t -> t -> bool
    val compare : t -> t -> int
    val of_bigint : Bigint.t -> t
    val to_bigint : t -> Bigint.t
    val ( + ) : t -> t -> t
    val ( +% ) : t -> Amount.t -> t
    val ( - ) : t -> t -> t
    val ( -% ) : t -> Amount.t -> t
    val ( = ) : t -> t -> bool
    val neg : t -> t
    val ( ~$ ) : int -> t
    val sexp_of_t : t -> Sexp.t

    val of_amount :
      Amount.t -> Debit_credit.t -> on_debit:[ `Incr | `Decr ] -> t
  end = struct
    type t = Bigint.t [@@deriving sexp_of]

    let equal = Bigint.equal
    let compare = Bigint.compare
    let of_bigint x = x
    let to_bigint x = x
    let ( + ) x y = Bigint.(x + y)
    let ( +% ) x y = x + of_bigint (Amount.to_bigint y)
    let ( - ) x y = Bigint.(x - y)
    let ( -% ) x y = x - of_bigint (Amount.to_bigint y)
    let ( = ) = equal
    let neg = Bigint.neg
    let ( ~$ ) = Fn.compose of_bigint Bigint.of_int

    let of_amount x (dc : Debit_credit.t) ~on_debit =
      match (dc, on_debit) with
      | Debit, `Incr -> of_bigint (Amount.to_bigint x)
      | Credit, `Incr -> neg (of_bigint (Amount.to_bigint x))
      | Debit, `Decr -> neg (of_bigint (Amount.to_bigint x))
      | Credit, `Decr -> of_bigint (Amount.to_bigint x)
  end
end

module Commodity_id = struct
  type t = string [@@deriving equal, compare, sexp]

  module Map = Map.Make (struct
    type nonrec t = t [@@deriving equal, compare, sexp]
  end)
end

(*
type scalar =
  | Amount of Money.Amount.t
  | Rate of { in_primary_commodity : Money.Amount.t; rate : Bigdecimal.t }
[@@deriving equal, compare, sexp_of] *)

module Gh = struct
  (* The category of the five top-level categories *)
  type global

  (* The five top-level categories *)
  type asset
  type equity
  type expense
  type income
  type liability

  (* No subcategories *)
  type final
end

module Account_structure0 (F : sig
  type 'a t
end) =
struct
  type 'a f =
    | Accounts_payable : Gh.final f F.t -> Gh.liability f
    | Accounts_receivable : Gh.final f F.t -> Gh.asset f
    | Bank : Gh.final f F.t -> Gh.asset f
    | Cash : Gh.final f F.t -> Gh.asset f
    | Credit : Gh.final f F.t -> Gh.liability f
    | Mutual_fund : Gh.final f F.t -> Gh.asset f
    | Stock : Gh.final f F.t -> Gh.asset f

  type t0 =
    | Asset of Gh.asset f F.t
    | Equity of Gh.equity f F.t
    | Expense of Gh.expense f F.t
    | Income of Gh.income f F.t
    | Liability of Gh.liability f F.t
end

module Account_type = struct
  type 'a elem = Leaf | Node of 'a

  include Account_structure0 (struct
    type 'a t = 'a elem
  end)
end

module Account_structure (F : sig
  type 'a t
end) =
struct
  include Account_structure0 (F)

  module Mapper = struct
    type nonrec 'b t = {
      car :
        'a.
        'a f F.t ->
        ('a Account_type.f Account_type.elem -> Account_type.t0) ->
        ('b * 'a f F.t) option;
    }

    let map (type b c) (f : c t) (mkt : b Account_type.f -> Account_type.t0) :
        b f -> (c * b f) option =
      let open Option.Let_syntax in
      function
      | Accounts_payable v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Accounts_payable el))
          in
          (c, Accounts_payable v')
      | Accounts_receivable v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Accounts_receivable el))
          in
          (c, Accounts_receivable v')
      | Bank v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Bank el)) in
          (c, Bank v')
      | Cash v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Cash el)) in
          (c, Cash v')
      | Credit v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Credit el)) in
          (c, Credit v')
      | Mutual_fund v ->
          let%map c, v' =
            f.car v (fun el -> mkt (Account_type.Mutual_fund el))
          in
          (c, Mutual_fund v')
      | Stock v ->
          let%map c, v' = f.car v (fun el -> mkt (Account_type.Stock el)) in
          (c, Stock v')
  end
end

module Typed_account_path = struct
  type 'a elem = Leaf | Node of string * 'a | Ind of string * 'a elem

  include Account_structure (struct
    type 'a t = 'a elem
  end)
end

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

  module Map = Map.Make (struct
    type nonrec t = t [@@deriving compare, sexp]
  end)
end

module Account_hierarchy = struct
  (* The contents of an account of category 'a *)
  type 'a core =
    (* Comprises of subaccounts of its subcategories *)
    | Node of 'a String.Map.t
    (* Comprises of subaccounts of its own category *)
    | Ind of 'a account String.Map.t
    (* Has no subaccounts, has a balance in a certain commodity *)
    | Leaf of Commodity_id.t * Money.Diff.t

  and extra = { description : String.t }
  and 'a account = extra * 'a core

  include Account_structure (struct
    type 'a t = 'a account
  end)

  (* All accounts *)
  type world = t0 String.Map.t

  let rec alter_aux (subaid : Account_path.t)
      (f :
        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
      : 'a Mapper.t =
    {
      car =
        (fun in_acc mkt ->
          let open Option.Let_syntax in
          match (subaid, in_acc) with
          | [], (extra, Leaf (acc_comm, acc_bal)) ->
              let x, acc_bal' = f (mkt Account_type.Leaf) acc_comm acc_bal in
              Some (x, (extra, Leaf (acc_comm, acc_bal')))
          | [], _ -> None
          | subaid0 :: subaid, (extra, Node subaccs) ->
              let open Option.Let_syntax in
              let%bind subacc = Map.find subaccs subaid0 in
              let%map x, subacc' =
                Mapper.map (alter_aux subaid f) (fun k -> mkt (Node k)) subacc
              in
              (x, (extra, Node (Map.set subaccs ~key:subaid0 ~data:subacc')))
          | subaid0 :: subaid, (extra, Ind subaccs) ->
              let open Option.Let_syntax in
              let%bind subacc = Map.find subaccs subaid0 in
              let%map x, subacc' = (alter_aux subaid f).car subacc mkt in
              (x, (extra, Ind (Map.set subaccs ~key:subaid0 ~data:subacc')))
          | _ :: _, (_, Leaf _) -> None);
    }

  let alter (aid : Account_path.t)
      (f :
        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
      (w : world) : ('a * world) option =
    match aid with
    | [] -> None
    | aid0 :: subaid -> (
        let open Option.Let_syntax in
        match%bind Map.find w aid0 with
        | Asset acc ->
            let%map x, acc' = (alter_aux subaid f).car acc (fun k -> Asset k) in
            (x, Map.set w ~key:aid0 ~data:(Asset acc'))
        | Expense acc ->
            let%map x, acc' =
              (alter_aux subaid f).car acc (fun k -> Expense k)
            in
            (x, Map.set w ~key:aid0 ~data:(Expense acc'))
        | Income acc ->
            let%map x, acc' =
              (alter_aux subaid f).car acc (fun k -> Income k)
            in
            (x, Map.set w ~key:aid0 ~data:(Income acc'))
        | Liability acc ->
            let%map x, acc' =
              (alter_aux subaid f).car acc (fun k -> Liability k)
            in
            (x, Map.set w ~key:aid0 ~data:(Liability acc'))
        | Equity acc ->
            let%map x, acc' =
              (alter_aux subaid f).car acc (fun k -> Equity k)
            in
            (x, Map.set w ~key:aid0 ~data:(Equity acc')))

  let world_inst : world =
    String.Map.of_alist_exn
      [
        ( "Assets",
          Asset
            ( { description = "assets" },
              Ind
                (String.Map.of_alist_exn
                   [
                     ( "Current",
                       ( { description = "current" },
                         Node
                           (String.Map.of_alist_exn
                              [
                                ( "Checking",
                                  Bank
                                    ( { description = "bnak accounts" },
                                      Ind
                                        (String.Map.of_alist_exn
                                           [
                                             ( "ING",
                                               ( { description = "ING bank" },
                                                 Leaf ("EUC", Money.Diff.(~$0))
                                               ) );
                                             ( "N26",
                                               ( { description = "ING bank" },
                                                 Leaf ("EUC", Money.Diff.(~$0))
                                               ) );
                                           ]) ) );
                              ]) ) );
                   ]) ) );
      ]
end

type bal_assert = {
  account : Account_path.t;
  labels : Labels.t;
  bal : Money.Diff.t;
}
[@@deriving sexp_of]

module Tx : sig
  type entry = {
    dc : Debit_credit.t;
    commodity : Commodity_id.t;
    amount : Money.Amount.t;
    assertion : Money.Diff.t option;
  }

  (* Private because we only want to allow constructing balanced transactions. *)
  type t = private {
    cleared : Date.t option;
    entries : entry Account_path.Map.t;
    labels : Labels.t;
  }

  type error = Unbalanced

  val make :
    cleared:Date.t option ->
    entries:entry Account_path.Map.t ->
    labels:Labels.t ->
    (t, error) result

  val sexp_of_t : t -> Sexp.t
end = struct
  type entry = {
    dc : Debit_credit.t;
    commodity : Commodity_id.t;
    amount : Money.Amount.t;
    assertion : Money.Diff.t option;
  }
  [@@deriving sexp_of]

  type t = {
    cleared : Date.t option;
    entries : entry Account_path.Map.t;
    labels : Labels.t;
  }
  [@@deriving sexp_of]

  type error = Unbalanced

  let is_balanced entries =
    Map.fold entries ~init:Commodity_id.Map.empty
      ~f:(fun ~key:_ ~data comm_balances ->
        Map.update comm_balances data.commodity ~f:(fun ocomm_bal ->
            let comm_bal = Option.value ocomm_bal ~default:Money.Diff.(~$0) in
            match data.dc with
            | Debit_credit.Debit -> Money.Diff.(comm_bal +% data.amount)
            | Debit_credit.Credit -> Money.Diff.(comm_bal -% data.amount)))
    |> Map.for_all ~f:(fun comm_bal -> Money.Diff.(comm_bal = ~$0))

  let make ~cleared ~entries ~labels =
    if not (is_balanced entries) then Error Unbalanced
    else Ok { cleared; entries; labels }
end

type item =
  | Tx_item of Tx.t
  | Bal_assert_item of bal_assert (*| Account_decl_item of account_decl*)
[@@deriving sexp_of]

type t = item list [@@deriving sexp_of]

module World = struct
  type t = Account_hierarchy.world

  let empty : t = String.Map.empty

  (** Update the balance (debit/credit ([dc])) of account [aid] [by_amount]
      (commodity: [in_comm]) in [world], giving the updated world and the pre
      and post balances for [aid] iff the account exists in [world]. *)
  let update_bal aid dc by_amount in_comm (world : t) :
      (Money.Diff.t * Money.Diff.t * t) option =
    let open Option.Let_syntax in
    let%bind mres, world' =
      Account_hierarchy.alter aid
        (fun acc_type acc_comm acc_bal ->
          if not ([%equal: Commodity_id.t] acc_comm in_comm) then (None, acc_bal)
          else
            match acc_type with
            | Asset _ ->
                let acc_bal' =
                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Incr)
                in
                (Some (acc_bal, acc_bal'), acc_bal')
            | Expense _ ->
                let acc_bal' =
                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Incr)
                in
                (Some (acc_bal, acc_bal'), acc_bal')
            | Income _ ->
                let acc_bal' =
                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
                in
                (Some (acc_bal, acc_bal'), acc_bal')
            | Liability _ ->
                let acc_bal' =
                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
                in
                (Some (acc_bal, acc_bal'), acc_bal')
            | Equity _ ->
                let acc_bal' =
                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
                in
                (Some (acc_bal, acc_bal'), acc_bal'))
        world
    in
    let%map pre_bal, post_bal = mres in
    (pre_bal, post_bal, world')

  let apply_tx (tx : Tx.t) world : t option =
    Map.fold_option tx.entries ~init:world
      ~f:(fun ~key:aid ~(data : Tx.entry) world ->
        let open Option.Let_syntax in
        let%bind _old_bal, new_bal, world =
          update_bal aid data.dc data.amount data.commodity world
        in
        match data.assertion with
        | None -> Some world
        | Some bal_ass ->
            if Money.Diff.(bal_ass = new_bal) then Some world else None)

  let apply : item -> t -> t option = function
    | Tx_item tx -> apply_tx tx
    | Bal_assert_item ba -> apply_ba ba
end

module Ctxd_item = struct end

let make = Fn.id