From Coq Require Import ZArith_base.
From Coq Require Import List. Import ListNotations.
From ConCert.Execution Require Import Monad.
From ConCert.Execution Require Import ResultMonad.
From ConCert.Execution Require Import Serializable.
From ConCert.Execution Require Import Blockchain.
From ConCert.Execution Require Import Containers.
From ConCert.Execution Require Import ContractCommon.
From ConCert.Utils Require Import Extras.
From ConCert.Utils Require Import RecordUpdate.

Section EIP20Token.
  Context {BaseTypes : ChainBase}.
  Set Nonrecursive Elimination Schemes.

  Definition TokenValue := N.
  Open Scope N_scope.
  Open Scope bool.

  Inductive Msg :=
  | transfer : Address -> TokenValue -> Msg
  | transfer_from : Address -> Address -> TokenValue -> Msg
  | approve : Address -> TokenValue -> Msg.

  Record State :=
    build_state {
        total_supply : TokenValue;
        balances : AddressMap.AddrMap TokenValue;
        allowances : AddressMap.AddrMap (AddressMap.AddrMap TokenValue)
    }.

  Record Setup :=
    build_setup {
      owner : Address;
      init_amount : TokenValue;
    }.

  Definition Error : Type := nat.
  Definition default_error : Error := 1%nat.


  Definition error {T : Type} : result T Error :=
    Err default_error.

  Definition init (chain : Chain)
                  (ctx : ContractCallContext)
                  (setup : Setup)
                  : result State Error :=
    Ok {| total_supply := setup.(init_amount);
            balances := AddressMap.add setup.(owner)
                                       setup.(init_amount)
                                       AddressMap.empty;
            allowances := AddressMap.empty
      |}.

  Definition increment_balance (m : AddressMap.AddrMap TokenValue)
                               (addr : Address)
                               (inc : TokenValue)
                               : AddressMap.AddrMap TokenValue :=
    match AddressMap.find addr m with
    | Some old => AddressMap.add addr (old + inc) m
    | None => AddressMap.add addr inc m
    end.

  Definition try_transfer (from : Address)
                          (to : Address)
                          (amount : TokenValue)
                          (state : State)
                          : result State Error :=
    let from_balance := with_default 0 (AddressMap.find from state.(balances)) in
    if from_balance <? amount
    then error
    else let new_balances := AddressMap.add from (from_balance - amount) state.(balances) in
         let new_balances := increment_balance new_balances to amount in
         Ok (state<|balances := new_balances|>).

  Definition try_transfer_from (delegate : Address)
                               (from : Address)
                               (to : Address)
                               (amount : TokenValue)
                               (state : State)
                               : result State Error :=
  do from_allowances_map <- result_of_option (AddressMap.find from state.(allowances)) default_error ;
  do delegate_allowance <- result_of_option (AddressMap.find delegate from_allowances_map) default_error ;
  let from_balance := with_default 0 (AddressMap.find from state.(balances)) in
  if (delegate_allowance <? amount) || (from_balance <? amount)
  then error
  else let new_allowances := AddressMap.add delegate (delegate_allowance - amount) from_allowances_map in
       let new_balances := AddressMap.add from (from_balance - amount) state.(balances) in
       let new_balances := increment_balance new_balances to amount in
       Ok (state<|balances := new_balances|><|allowances ::= AddressMap.add from new_allowances|>).

  Definition try_approve (caller : Address)
                         (delegate : Address)
                         (amount : TokenValue)
                         (state : State)
                         : result State Error :=
    match AddressMap.find caller state.(allowances) with
    | Some caller_allowances =>
      Ok (state<|allowances ::= AddressMap.add caller (AddressMap.add delegate amount caller_allowances)|>)
    | None =>
      Ok (state<|allowances ::= AddressMap.add caller (AddressMap.add delegate amount AddressMap.empty)|>)
    end.

  Open Scope Z_scope.
  Definition receive (chain : Chain)
                     (ctx : ContractCallContext)
                     (state : State)
                     (maybe_msg : option Msg)
                     : result (State * list ActionBody) Error :=
    let sender := ctx.(ctx_from) in
    

    if ctx.(ctx_amount) >? 0
    then error
    else
      match maybe_msg with
      | Some (transfer to amount) =>
          without_actions (try_transfer sender to amount state)
      | Some (transfer_from from to amount) =>
          without_actions (try_transfer_from sender from to amount state)
      | Some (approve delegate amount) =>
          without_actions (try_approve sender delegate amount state)
      

      | None => error
      end.
  Close Scope Z_scope.

  Definition contract : Contract Setup Msg State Error :=
    build_contract init receive.

  (* sum of all balances in the token state *)
  Definition sum_balances (state : EIP20Token.State) :=
    sumN (fun '(k, v) => v) (FMap.elements (balances state)).

  Definition get_allowance state from delegate :=
    with_default 0 (FMap.find delegate (with_default
      (@FMap.empty (FMap Address TokenValue) _) (FMap.find from (allowances state)))).

End EIP20Token.