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 Account = 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 t 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 t = extra * 'a core

  (* 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

  (* Subaccounts under the five top-level categories *)
  type 'a f =
    | Accounts_payable : final f t -> liability f
    | Accounts_receivable : final f t -> asset f
    | Bank : final f t -> asset f
    | Cash : final f t -> asset f
    | Credit : final f t -> liability f
    | Mutual_fund : final f t -> asset f
    | Stock : final f t -> asset f

  module Ft_mapper = struct
    type nonrec 'b t = { car : 'a. 'a f t -> ('b * 'a f t) option }

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

  module Top_level (F : sig
    type 'a t
  end) =
  struct
    type t =
      | Asset of asset F.t
      | Equity of equity F.t
      | Expense of expense F.t
      | Income of income F.t
      | Liability of liability F.t
  end

  module Top_level_type = Top_level (struct
    type nonrec 'a t = unit
  end)

  (* I swear the bullshit stops here *)
  module F0 = struct
    include Top_level (struct
      type nonrec 'a t = 'a f t
    end)

    let type_ : t -> Top_level_type.t = function
      | Asset _ -> Asset ()
      | Equity _ -> Equity ()
      | Expense _ -> Expense ()
      | Income _ -> Income ()
      | Liability _ -> Liability ()
  end

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

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

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

  let world_inst : world =
    String.Map.of_alist_exn
      [
        ( "Assets",
          F0.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.world

  let empty : t = String.Map.empty

  let update_bal_fn = (f :
        Account.Top_level_type.t ->
        Commodity_id.t ->
        Money.Diff.t ->
        ('a * Money.Diff.t) option)
  
  (* Stretching the type system a little :) *)
  let rec update_bal_aux ttype subaid
      (f :
        Account.Top_level_type.t ->
        Commodity_id.t ->
        Money.Diff.t ->
        ('a * Money.Diff.t) option) : 'a Account.Ft_mapper.t =
    {
      car =
        (fun in_acc ->
          let open Option.Let_syntax in
          match (subaid, in_acc) with
          | [], (extra, Account.Leaf (acc_comm, acc_bal)) ->
              let%bind x, acc_bal' = f ttype acc_comm acc_bal in
              Some (x, (extra, Account.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' =
                Account.Ft_mapper.map (update_bal_aux ttype subaid f) subacc
              in
              ( x,
                ( extra,
                  Account.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' = (update_bal_aux ttype subaid f).car subacc in
              ( x,
                (extra, Account.Ind (Map.set subaccs ~key:subaid0 ~data:subacc'))
              )
          | _ :: _, (_, Leaf _) -> None);
    }

  (** 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 =
    match aid with
    | [] -> None
    | aid0 :: subaid -> (
        let open Option.Let_syntax in
        match%bind Map.find world aid0 with
        | Asset acc ->
            let diff = Money.Diff.of_amount by_amount dc ~on_debit:`Incr in
            let%map bals, acc' = (update_bal_aux subaid diff in_comm).car acc in
            (bals, Map.set world ~key:aid0 ~data:(Asset acc'))
        | Expense acc ->
            let diff = Money.Diff.of_amount by_amount dc ~on_debit:`Incr in
            let%map bals, acc' = (update_bal_aux subaid diff in_comm).car acc in
            (bals, Map.set world ~key:aid0 ~data:(Expense acc'))
        | Income acc ->
            let diff = Money.Diff.of_amount by_amount dc ~on_debit:`Decr in
            let%map bals, acc' = (update_bal_aux subaid diff in_comm).car acc in
            (bals, Map.set world ~key:aid0 ~data:(Income acc'))
        | Liability acc ->
            let diff = Money.Diff.of_amount by_amount dc ~on_debit:`Decr in
            let%map bals, acc' = (update_bal_aux subaid diff in_comm).car acc in
            (bals, Map.set world ~key:aid0 ~data:(Liability acc'))
        | Equity acc ->
            let diff = Money.Diff.of_amount by_amount dc ~on_debit:`Decr in
            let%map bals, acc' = (update_bal_aux subaid diff in_comm).car acc in
            (bals, Map.set world ~key:aid0 ~data:(Equity acc')))

  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