module Type = struct
  type t =
    | Int
    | Float
    | String

  let to_string = function
    | Int -> "int"
    | Float -> "float"
    | String -> "string"

  let merge a b =
    match a, b with
    | Int, Float -> Float
    | Float, Int -> Float
    | a, b when a = b -> a
    | _ -> failwith "Type.merge"

end

exception Invalid_type of Type.t

module Value = struct
  type t =
    | Int of int
    | Float of float
    | String of string
    | Nop (* return of system operations *)

  let to_string = function
    | Int n -> string_of_int n
    | Float n -> string_of_float n
    | String s -> s
    | Nop -> "nop"

  let of_token = function
    | Token.Int n -> Int n
    | Float n -> Float n
    | _ -> invalid_arg "Value.of_token"

  let typeof = function
    | Int _ -> Type.Int
    | Float _ -> Type.Float
    | String _ -> Type.String
    | Nop -> failwith "Value.typeof"

  let promote = function
    | Int n -> Float (float n)
    | Float n -> Float n
    | _ -> failwith "Value.promote"
end

(* binary operator *)
(* binary operator has type 'a -> 'a -> 'b. *)
module Binop = struct
  type t =
    | Add | Sub | Mul | Div (* arithmetics *)
    | Mod (* modular operation *)
    | Exp (* exponentation *)

  let to_string = function
    | Add -> "+"
    | Sub -> "-"
    | Mul -> "*"
    | Div -> "/"
    | Mod -> "%"
    | Exp -> "^"

  let vi f a b =
    let open Value in
    match a, b with
    | Int a, Int b -> Int (f a b)
    | _ -> raise @@ Invalid_type Int

  let vf f a b =
    let open Value in
    match a, b with
    | Float a, Float b -> Float (f a b)
    | _ -> raise @@ Invalid_type Float

  let operators =
    let open Type in
    [
      Add, [Int, vi Int.add; Float, vf Float.add];
      Sub, [Int, vi Int.sub; Float, vf Float.sub];
      Mul, [Int, vi Int.mul; Float, vf Float.mul];
      Div, [Int, vi Int.div; Float, vf Float.div];
      Mod, [Int, vi Int.rem; Float, vf Float.rem];
      Exp, [Float, vf Float.pow];
    ]
    |> List.to_seq
    |> Hashtbl.of_seq

  let get op typ =
    Hashtbl.find operators op
    |> List.assoc_opt typ
end

type t =
  | Value of Value.t
  | Var of string
  | Let of string * t
  | Binop of t * Binop.t * t
  | Set_binop_pre of Binop.t * t
  | Get_binop_pre of Binop.t
  | Set_binop_aso of Binop.t * string
  | Get_binop_aso of Binop.t

let value v = Value v

let binop left op right =
  Binop (left, op, right)

(* print ast LISP style. *)
let print ast =
  let pr = Printf.printf in
  let pv v = pr "%s" @@ Value.to_string v in
  let rec aux = function
    | Value n -> pv n
    | Var v -> pr "%s" v
    | Let (v, e) ->
      pr "(let %s " v;
      aux e;
      pr ")"
    | Binop (left, op, right) -> begin
        let op = Binop.to_string op in
        pr "(%s " op; aux left; pr " "; aux right; pr ")";
      end
    | Set_binop_pre (op, pre) ->
      pr "(set_pre %s " (Binop.to_string op);
      aux pre;
      pr ")"
    | Get_binop_pre op ->
      pr "(get_pre %s)" (Binop.to_string op)
    | Set_binop_aso (op, aso) ->
      pr "(set_assoc %s %s)" (Binop.to_string op) aso
    | Get_binop_aso op ->
      pr "(get_pre %s)" (Binop.to_string op)
  in
  aux ast; pr "\n"