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_node : expr_node) :
    (lox_value, runtime_error) result =
  let { pos; expr } = expr_node in
  match expr with
  | Literal literal -> value_of_literal literal |> Result.ok
  | 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
          RuntimeError.make pos msg |> Result.error)
  | Assignment { name; expr } ->
      if not (Env.is_defined env name) then
        let msg = Printf.sprintf "tried to assign to undefined variable %s" name in
        RuntimeError.make pos msg |> Result.error
      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 -> Number (-.x) |> Result.ok
      | Not, value -> Bool (lox_value_to_bool value |> not) |> Result.ok
      | _, _ ->
          let msg =
            Printf.sprintf "Invalid operant of type %s to operator %s"
              (type_string_of_lox_value expr) (show_unary_op op)
          in
          RuntimeError.make pos msg |> Result.error)
  | 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 -> String (a ^ b) |> Result.ok
      | Number x, Plus, Number y -> Number (x +. y) |> Result.ok
      | Number x, Minus, Number y -> Number (x -. y) |> Result.ok
      | Number x, Mul, Number y -> Number (x *. y) |> Result.ok
      | Number x, Div, Number y ->
          if y <> 0. then Number (x /. y) |> Result.ok
          else
            let msg = "Division by 0" in
            RuntimeError.make pos msg |> Result.error
      | Bool b, And, Bool c -> Bool (b && c) |> Result.ok
      | Bool b, Or, Bool c -> Bool (b || c) |> Result.ok
      | _, Equal, _ -> Ok (Bool (left = right))
      | Number x, Greater, Number y -> Bool (x > y) |> Result.ok
      | Number x, GreaterEqual, Number y -> Bool (x >= y) |> Result.ok
      | Number x, Less, Number y -> Bool (x < y) |> Result.ok
      | Number x, LessEqual, Number y -> Bool (x <= y) |> Result.ok
      | String a, Greater, String b -> Bool (a > b) |> Result.ok
      | String a, GreaterEqual, String b -> Bool (a >= b) |> Result.ok
      | String a, Less, String b -> Bool (a < b) |> Result.ok
      | String a, LessEqual, String b -> Bool (a <= b) |> Result.ok
      | _, _, _ ->
          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
          RuntimeError.make pos msg |> Result.error)
  | 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)
  | Call { callee; args } -> (
      let* callee = interpret_expr env callee in
      let f (acc : (lox_value list, runtime_error) result) (arg : expr_node) :
          (lox_value list, runtime_error) result =
        match acc with
        | Ok acc ->
            let* arg = interpret_expr env arg in
            Ok (arg :: acc)
        | Error e -> Error e
      in
      let* args = List.fold_left f (Ok []) args in
      let args = List.rev args in
      match callee with
      | Function (NativeFunction { name; arity; fn }) ->
          let args_len = List.length args in
          if args_len <> arity then
            let msg =
              Printf.sprintf "Native Function %s has arity %d, but was called with %d args" name
                args_len arity
            in
            RuntimeError.make pos msg |> Result.error
          else fn args |> Result.map_error (RuntimeError.make pos)
      | Function (LoxFunction { name; arity; arg_names; body }) ->
          let args_len = List.length args in
          if args_len <> arity then
            let msg =
              Printf.sprintf "Function %s has arity %d, but was called with %d args" name args_len
                arity
            in
            RuntimeError.make pos msg |> Result.error
          else
            let env = Env.push_frame env in
            let () =
              List.iter2 (fun name value -> assert (Env.define env name value)) arg_names args
            in
            let* () = interpret_stmt env body in
            Ok Nil
      | _ ->
          ignore args;
          let msg = Printf.sprintf "%s object is not callable" (type_string_of_lox_value callee) in
          RuntimeError.make pos msg |> Result.error)

and interpret_stmt (env : environment) (stmt_node : stmt_node) : (unit, runtime_error) result =
  let { pos; stmt } = stmt_node in
  ignore pos;
  match stmt with
  | Expr expr ->
      let* value = interpret_expr env expr in
      ignore value;
      Ok ()
  | Break -> RuntimeError.break () |> Result.error
  | Continue -> RuntimeError.continue () |> Result.error
  | 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 a variable of that name was already defined" name
        in
        RuntimeError.make pos msg |> Result.error
  | FunDecl { name; arg_names; body } ->
      let fn = make_lox_function name arg_names body in
      let success = Env.define env name fn in
      if success then Ok ()
      else
        let msg =
          Printf.sprintf
            "Tried to define function %s, but a variable of that name was already defined" name
        in
        RuntimeError.make pos msg |> Result.error
  | 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 ())
  | While { cond; body } ->
      let* cond = interpret_expr env cond in
      let cond = lox_value_to_bool cond in
      if cond then
        let result = interpret_stmt env body in
        match result with
        | Ok () | Error Continue -> interpret_stmt env stmt_node
        | Error Break -> Ok ()
        | Error e -> Error e
      else Ok ()
  | For { init; cond; update; body } ->
      let env = Env.enter env in
      let* () = init |> Option.map (interpret_stmt env) |> Option.value ~default:(Ok ()) in
      let eval_cond () =
        cond
        |> Option.map (interpret_expr env)
        |> Option.value ~default:(Ok (Value.Bool true))
        |> Result.map Value.lox_value_to_bool
      in
      let do_update () =
        update |> Option.map (interpret_expr env) |> Option.value ~default:(Ok Value.Nil)
      in
      let rec loop () =
        let* cond = eval_cond () in
        if cond then
          let result = interpret_stmt env body in
          match result with
          | Ok () | Error Continue ->
              let* value = do_update () in
              ignore value;
              loop ()
          | Error Break -> Ok ()
          | Error e -> Error e
        else Ok ()
      in
      loop ()