module VariableBindingMap = Map.Make(String)


type value_type = 
| Int of int
| Fun of function_type
and scope = {
  parent: scope option;
  bindings: value_type VariableBindingMap.t;
}
and function_type = {
  name: string;
  body: Parser.expr_tree;
  scope: scope;
}

let rec eval_expr (scope: scope) (expr: Parser.expr_tree): value_type = 
  match expr with
  | Parser.LetExpr (Parser.Let (name, value_expr, in_expr)) -> 
    eval_let_expr scope name value_expr in_expr
  | Parser.FunExpr (Parser.Fun (name, body_expr)) -> 
    eval_fun_expr scope name body_expr
  | Parser.IfExpr (Parser.If (cond_expr, then_expr, else_expr)) -> 
    eval_if_expr scope cond_expr then_expr else_expr
  | Parser.BinOpExpr (op, left_expr, right_expr) -> 
    eval_bin_op_expr scope op left_expr right_expr
  | Parser.MonoOpExpr (_op, _expr) -> 
    failwith "Not implemented"
  | Parser.Identifier(name) ->
    let rec find_binding scope = 
      match scope with
      | None -> failwith "Unbound variable"
      | Some s -> 
        match VariableBindingMap.find_opt name s.bindings with
        | Some v -> v
        | None -> find_binding s.parent in
    find_binding (Some scope)
  | Parser.Number(n) -> Int n
and eval_if_expr scope cond_expr then_expr else_expr = 
  let cond = eval_expr scope cond_expr in
  (match cond with
  | Int 0 -> eval_expr scope else_expr
  | _ -> eval_expr scope then_expr)
and eval_let_expr scope name value_expr in_expr =
  let value = eval_expr scope value_expr in
  let new_scope = { scope with bindings = VariableBindingMap.add name value scope.bindings } in
  eval_expr new_scope in_expr
and eval_fun_expr scope name body_expr =
  Fun { name = name; body = body_expr; scope = scope }
and eval_bin_op_expr scope op left_expr right_expr = 
  let left = eval_expr scope left_expr in
    let right = eval_expr scope right_expr in
    (match op with
    | Add -> (
      match (left, right) with
      | (Int l, Int r) -> Int (l + r)
      | _ -> failwith "Type error"
    )
    | Sub -> (
      match (left, right) with
      | (Int l, Int r) -> Int (l - r)
      | _ -> failwith "Type error"
    )
    | Mul -> (
      match (left, right) with
      | (Int l, Int r) -> Int (l * r)
      | _ -> failwith "Type error"
    )
    | Div -> (
      match (left, right) with
      | (Int l, Int r) -> Int (l / r)
      | _ -> failwith "Type error"
    )
    | Mod -> (
      match (left, right) with
      | (Int l, Int r) -> Int (l mod r)
      | _ -> failwith "Type error"
    )
    | Pow -> (
      match (left, right) with
      | (Int l, Int r) -> Int (int_of_float (float_of_int l ** float_of_int r))
      | _ -> failwith "Type error"
    ))


let eval_str (str: string): value_type = 
  let tokens = Lexer.lex_tokens_seq str in
  let tokens = tokens |> Seq.map (fun (x,_) -> x) in
  let expr = Parser.get_expr_tree_from_tokens tokens in
  match expr with
  | Some e -> eval_expr { parent = None; bindings = VariableBindingMap.empty } e
  | None -> failwith "Parse error"


let%test "test eval_str 1" =
  let result = eval_str "let x = 1 in x" in
  match result with
  | Int n -> n = 1
  | _ -> false