let ( let* ) = Result.bind

open Environment
open Error
open Expr
open Stmt
open Value

let value_of_literal (literal : literal) : Value.lox_value =
  match literal with String s -> String s | Number x -> Number x | Bool b -> Bool b | Nil -> Nil

let rec interpret_expr (env : environment) (expr : expr_node) : (lox_value, runtime_error) result =
  let { pos; expr } = expr in
  match expr with
  | Literal literal -> Ok (value_of_literal literal)
  | Variable name -> (
      let value_opt = Env.get env name in
      match value_opt with
      | Some x -> Ok x
      | None ->
          let msg = Printf.sprintf "name \"%s\" is not defined" name in
          Error (RuntimeError.make pos msg))
  | Assignment { name; expr } ->
      if not (Env.is_defined env name) then
        let msg = Printf.sprintf "tried to assign to undefined variable %s" name in
        Error (RuntimeError.make pos msg)
      else
        let* value = interpret_expr env expr in
        Env.update env name value;
        Ok value
  | Unary { op; expr } -> (
      let* expr = interpret_expr env expr in
      match (op, expr) with
      | Neg, Number x -> Ok (Number (-.x))
      | Not, Bool b -> Ok (Bool (not b))
      | _, _ ->
          let msg =
            Printf.sprintf "Invalid operant of type %s to operator %s"
              (type_string_of_lox_value expr) (show_unary_op op)
          in
          Error (RuntimeError.make pos msg))
  | Binary { op; left; right } -> (
      let* left = interpret_expr env left in
      let* right = interpret_expr env right in
      match (left, op, right) with
      | String a, Plus, String b -> Ok (String (a ^ b))
      | Number x, Plus, Number y -> Ok (Number (x +. y))
      | Number x, Minus, Number y -> Ok (Number (x -. y))
      | Number x, Mul, Number y -> Ok (Number (x *. y))
      | Number x, Div, Number y ->
          if y <> 0. then Ok (Number (x /. y))
          else
            let msg = "Division by 0" in
            Error { pos; msg }
      | Bool b, And, Bool c -> Ok (Bool (b && c))
      | Bool b, Or, Bool c -> Ok (Bool (b || c))
      | _, Equal, _ -> Ok (Bool (left = right))
      | Number x, Greater, Number y -> Ok (Bool (x > y))
      | Number x, GreaterEqual, Number y -> Ok (Bool (x >= y))
      | Number x, Less, Number y -> Ok (Bool (x < y))
      | Number x, LessEqual, Number y -> Ok (Bool (x <= y))
      | String a, Greater, String b -> Ok (Bool (a > b))
      | String a, GreaterEqual, String b -> Ok (Bool (a >= b))
      | String a, Less, String b -> Ok (Bool (a < b))
      | String a, LessEqual, String b -> Ok (Bool (a <= b))
      | _, _, _ ->
          let msg =
            Printf.sprintf "Invalid operands of type %s and %s to operator %s"
              (type_string_of_lox_value left) (type_string_of_lox_value right) (show_binary_op op)
          in
          Error { pos; msg })
  | Logical { op; left; right } -> (
      let* left = interpret_expr env left in
      match (op, lox_value_to_bool left) with
      | And, false | Or, true -> Ok left (* short circuit *)
      | _ -> interpret_expr env right)

let rec interpret_stmt (env : environment) (stmt : stmt_node) : (unit, runtime_error) result =
  let { pos; stmt } = stmt in
  ignore pos;
  match stmt with
  | Expr expr ->
      let* _ = interpret_expr env expr in
      Ok ()
  | Print expr ->
      let* value = interpret_expr env expr in
      print_endline (Value.string_of_lox_value value);
      Ok ()
  | VarDecl { name; init } ->
      let* init = Option.map (interpret_expr env) init |> Option.value ~default:(Ok Nil) in
      let success = Env.define env name init in
      if success then Ok ()
      else
        let msg = Printf.sprintf "Tried to define %s, but was already defined" name in
        Error (RuntimeError.make pos msg)
  | Block stmts ->
      let env = Env.enter env in
      let rec _interpret stmts =
        match stmts with
        | stmt :: tail ->
            let* _ = interpret_stmt env stmt in
            _interpret tail
        | [] -> Ok ()
      in
      _interpret stmts
  | If { cond; then_; else_ } ->
      let* cond = interpret_expr env cond in
      let cond = lox_value_to_bool cond in
      if cond then interpret_stmt env then_
      else Option.map (interpret_stmt env) else_ |> Option.value ~default:(Ok ())