open Ast

type command = [ `Parse | `DrawTree | `RunCBN | `RunCBV | `StepCBN | `StepCBV ]

let command_of_yojson = function
  | `String "Parse" -> Ok `Parse
  | `String "DrawTree" -> Ok `DrawTree
  | `String "RunCBN" -> Ok `RunCBN
  | `String "RunCBV" -> Ok `RunCBV
  | `String "StepCBN" -> Ok `StepCBN
  | `String "StepCBV" -> Ok `StepCBV
  | _ -> Error "invalid command"

type request = { expr : string; command : command }
[@@deriving of_yojson { exn = true }]

type response_expr = { expr : string; has_redex : bool; note : string option }
[@@deriving to_yojson { exn = true }]

type response =
  [ `Error of string
  | `Expr of response_expr
  | `Exprs of string option * string array
  | `Graphviz of string ]
[@@deriving to_yojson { exn = true }]

let%expect_test _ =
  Yojson.Safe.to_channel stdout (response_to_yojson (`Error "foo"));
  [%expect {| ["Error","foo"] |}]

let%expect_test _ =
  Yojson.Safe.to_channel stdout
    (response_to_yojson
       (`Expr { expr = "foo"; has_redex = false; note = None }));
  [%expect {| ["Expr",{"expr":"foo","has_redex":false,"note":null}] |}]

let%expect_test _ =
  Yojson.Safe.to_channel stdout
    (response_to_yojson (`Exprs (Some "A", [| "B"; "C" |])));
  [%expect {| ["Exprs","A",["B","C"]] |}]

let handle_request { expr; command } : response =
  try
    let buf = Lexing.from_string expr in
    let expr = Parse.expression buf in
    let expr = expr_of_parsetree expr in

    let expr_response ?(note : string option) (expr : expr ast) : response =
      `Expr
        {
          expr = show_expr expr;
          has_redex = Option.is_some (Eval.find_redex_cbn expr);
          note;
        }
    in

    let step_with (find_redex : expr ast -> expr index option) (expr : expr ast)
        : expr ast =
      match find_redex expr with
      | Some i -> Eval.reduce expr i
      | None -> failwith "no redex"
    in
    let run_with (find_redex : expr ast -> expr index option) :
        expr ast -> expr ast Seq.t =
      let rec loop (expr : expr ast) : expr ast Seq.t =
        Seq.cons expr (fun () ->
            match find_redex expr with
            | Some i -> loop (Eval.reduce expr i) ()
            | None -> Nil)
      in
      loop
    in

    match command with
    | `Parse -> expr_response expr
    | `StepCBN ->
        expr_response ~note:"After one CBN step:"
          (step_with Eval.find_redex_cbn expr)
    | `StepCBV ->
        expr_response ~note:"After one CBV step:"
          (step_with Eval.find_redex_cbv expr)
    | `DrawTree -> `Graphviz (Draw_tree.draw_tree expr)
    | `RunCBN ->
        `Exprs
          ( Some "Evaluating with CBN:",
            run_with Eval.find_redex_cbn expr
            |> Seq.take 100 |> Seq.map show_expr |> Array.of_seq )
    | `RunCBV ->
        `Exprs
          ( Some "Evaluating with CBV:",
            run_with Eval.find_redex_cbv expr
            |> Seq.take 100 |> Seq.map show_expr |> Array.of_seq )
  with
  | Ast.UnsupportedAst ast -> `Error ("unsupported AST: " ^ ast)
  | Eval.NotARedex expr -> `Error ("not a redex: " ^ show_expr expr)
  | Failure msg -> `Error msg
  | exn -> `Error ("uncaught exception: " ^ Printexc.to_string exn)