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

type 'a interpreter_result =
  | Ok of 'a
  | Break
  | Continue
  | Return of lox_value
  | Error of runtime_error

module InterpreterResult = struct
  type 'a t = 'a interpreter_result

  let ok value = Ok value
  let break () = Break
  let continue () = Continue
  let return (value : lox_value) = Return value
  let error e = Error e

  let map f result =
    match result with
    | Ok x -> f x |> ok
    | Break -> Break
    | Continue -> Continue
    | Return _ as r -> r
    | Error _ as e -> e

  let map_error f result =
    match result with
    | Error e -> f e |> error
    | Ok _ as x -> x
    | Break -> Break
    | Continue -> Continue
    | Return _ as r -> r

  let bind result f =
    match result with
    | Ok value -> f value
    | Break -> Break
    | Continue -> Continue
    | Return _ as r -> r
    | Error _ as e -> e
end

open InterpreterResult

let ( let* ) = InterpreterResult.bind

type expr_result = lox_value interpreter_result
type stmt_result = unit interpreter_result

let rec interpret_expr (env : environment) (expr_node : expr_node) : lox_value interpreter_result =
  let { pos; expr } = expr_node in
  match expr with
  | Literal literal -> value_of_literal literal |> 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 |> 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 |> 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) |> ok
      | Not, value -> Bool (lox_value_to_bool value |> not) |> 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 |> 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) |> ok
      | Number x, Plus, Number y -> Number (x +. y) |> ok
      | Number x, Minus, Number y -> Number (x -. y) |> ok
      | Number x, Mul, Number y -> Number (x *. y) |> ok
      | Number x, Div, Number y ->
          if y <> 0. then Number (x /. y) |> ok
          else
            let msg = "Division by 0" in
            RuntimeError.make pos msg |> error
      | Bool b, And, Bool c -> Bool (b && c) |> ok
      | Bool b, Or, Bool c -> Bool (b || c) |> ok
      | _, Equal, _ -> Ok (Bool (left = right))
      | Number x, Greater, Number y -> Bool (x > y) |> ok
      | Number x, GreaterEqual, Number y -> Bool (x >= y) |> ok
      | Number x, Less, Number y -> Bool (x < y) |> ok
      | Number x, LessEqual, Number y -> Bool (x <= y) |> ok
      | String a, Greater, String b -> Bool (a > b) |> ok
      | String a, GreaterEqual, String b -> Bool (a >= b) |> ok
      | String a, Less, String b -> Bool (a < b) |> ok
      | String a, LessEqual, String b -> Bool (a <= b) |> 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 |> 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 interpreter_result) arg =
        let* acc = acc in
        let* arg = interpret_expr env arg in
        Ok (arg :: acc)
      in
      let* args = List.fold_left f (Ok []) args in
      let args = List.rev args in
      match callee with
      | 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 |> error
          else
            match fn args with Ok value -> Ok value | Error s -> RuntimeError.make pos s |> error)
      | Function { name; env; 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 |> error
          else
            (* let env = Env.push_frame env in *)
            let env = Env.enter env in
            let () =
              List.iter2 (fun name value -> assert (Env.define env name value)) arg_names args
            in
            let result = interpret_stmt env body in
            match result with
            | Ok () -> Ok Nil
            | Return value -> Ok value
            | Error _ as e -> e
            | _ -> assert false)
      | _ ->
          ignore args;
          let msg = Printf.sprintf "%s object is not callable" (type_string_of_lox_value callee) in
          RuntimeError.make pos msg |> error)

and interpret_stmt (env : environment) (stmt_node : stmt_node) : unit interpreter_result =
  let { stmt; pos } = stmt_node in
  match stmt with
  | Expr expr ->
      let* value = interpret_expr env expr in
      ignore value;
      Ok ()
  | Break -> break ()
  | Continue -> continue ()
  | Return expr ->
      let* value = interpret_expr env expr in
      Return value
  | 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 |> error
  | FunDecl { name; arg_names; body } ->
      let fn = make_lox_function name env 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 |> 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 () | Continue -> interpret_stmt env stmt_node
        | Break -> Ok ()
        | other -> other
      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))
        |> 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 () | Continue ->
              let* value = do_update () in
              ignore value;
              loop ()
          | Break -> Ok ()
          | other -> other
        else Ok ()
      in
      loop ()