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

  let to_string = function
    | Int -> "int"
    | Float -> "float"
    | String -> "string"
    | Function -> "fun"
    | External -> "external"

  let supertype = function
    | Int -> Some Float
    | _ -> None
end

(* simple, untyped AST. *)
type t =
  | Value of value
  | Var of string
  | Let of string * t
  | Unary of operator * t
  | Binop of t * operator * t
  | Apply of t * t list (* function application *)
  | Set_binop_pre of operator * t
  | Get_binop_pre of operator
  | Set_binop_aso of operator * string
  | Get_binop_aso of operator

and value =
  | Int of int
  | Float of float
  | String of string
  | Function of string list * t
  | External of string
  | Nop (* return of system operations (will be deprecated) *)

and operator =
  | Add | Sub | Mul | Div (* arithmetics *)
  | Mod (* modular operation *)
  | Exp (* exponentation *)
  | Negate

exception Invalid_type of Type.t

module Value = struct
  type t = value

  let to_string = function
    | Int n -> string_of_int n
    | Float n -> string_of_float n
    | String s -> "\"" ^ s ^ "\""
    | Function (vars, _) ->
      Printf.sprintf "function with %d arguments" @@ List.length vars
    | External f -> "external " ^ f
    | Nop -> "nop"

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

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

(* operators *)
module Operator = struct
  type t = operator

  exception Unavailable of t

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

  let negate = function
    | Int n -> Int ~-n
    | Float n -> Float ~-.n
    | _ -> failwith "Operator.negate"

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

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

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

  let get_types op =
    match Hashtbl.find_opt operators op with
    | None -> raise @@ Unavailable op
    | Some p -> List.map fst p

  let get_unary = function
    | Negate -> negate
    | op -> raise @@ Unavailable op

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

let value v = Value v

let unary op t =
  Unary (op, t)

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 ")"
    | Unary (op, t) ->
      let op = Operator.to_string op in
      pr "(%s " op; aux t; pr ")"
    | Binop (left, op, right) ->
      let op = Operator.to_string op in
      pr "(%s " op; aux left; pr " "; aux right; pr ")"
    | Apply (f, args) ->
      pr "("; List.iter aux @@ f::args; pr ")"

    | Set_binop_pre (op, pre) ->
      pr "(set_pre %s " (Operator.to_string op);
      aux pre;
      pr ")"
    | Get_binop_pre op ->
      pr "(get_pre %s)" (Operator.to_string op)
    | Set_binop_aso (op, aso) ->
      pr "(set_assoc %s %s)" (Operator.to_string op) aso
    | Get_binop_aso op ->
      pr "(get_pre %s)" (Operator.to_string op)
  in
  aux ast; pr "\n"