let test = Testing.test

exception Misformated_Macro of Sexpr.t
exception Macro_Doesnt_App
exception Not_Valid_Macro_Arg_Syntax
exception FIXME of Sexpr.t

let num_expands = ref 0;;

let assoc_all k li =
  li
  |> List.find_all (fun (x, _) -> x = k)
  |> List.map (fun (_, v) -> v)

let () = test "SexprMacroExp.assoc_all" @@ begin fun () ->
  let li = [
    "asye","v4";
    "key","v1";
    "notkey","v2";
    "key","v3"
  ] in
  assert (assoc_all "key" li = ["v1";"v3"])
end

let rec eval (ctx: (string * (Sexpr.t * Sexpr.t)) list) s =
  let open Sexpr in
  let aggr_eval ctx s = begin
    let s, (ch,ctx) = aggr_flat_map (fun aggr s ->
      let ch1, ctx = aggr in
      let s, ch2, ctx = eval ctx s in
      s, ((ch1 || ch2), ctx)
    ) (false, ctx) s in
    s, ch, ctx
  end in
  let s, ch1, ctx = match s with
    Cons(Id ":macro", Cons(name, Cons(margs, rr))) ->
      let rr, ch, ctx = aggr_eval ctx rr in
      Cons(Id ":macro", Cons(name, Cons(margs, rr))), ch, ctx
  | s when is_list s -> begin
      aggr_eval ctx s
    end
  | _ ->
      s, false, ctx
  in
  let s, ch2, ctx = check_expands ctx s in
  s, (ch1 || ch2), ctx
(* expr -> list(new exprs) * bool:changed * ctx *)
and check_expands ctx s =
  let open Sexpr in
  match s with
    Cons(Id ":macro", r) -> begin
      let a = match r with
        Cons(Id(l), Cons(a,b)) -> l, (a,b)
      | _ -> raise @@ Misformated_Macro s
      in
      let ctx = (a :: ctx)
        |> List.sort (fun a b -> macro_req_score (fst (snd b)) - macro_req_score (fst (snd a)))
      in
      Nil, false, ctx
    end
  | Cons(Id ":atm-to-str", Cons(Id x, Nil))
  | Cons(Id ":atm-to-str", Cons(Str x, Nil)) ->
      li1 (Str x), true, ctx
  | Cons(Id ":atm-to-str", Cons(Int x, Nil)) ->
      li1 (Str (Int.to_string x)), true, ctx
  | Cons(Id ":atm-to-str", Cons(Flt x, Nil)) ->
      li1 (Str (Float.to_string x)), true, ctx
  | Cons(Id ":str-cat", Cons(Str l, Cons(Str r, Nil))) ->
      li1 (Str (l ^ r)), true, ctx
  | Cons(Id ":unpack", Cons(x, Nil)) when is_list x ->
      x, true, ctx
  | Cons(Id ":c-eq", Cons(l, Cons(r, Nil))) when l = r ->
      Nil, true, ctx
  | Cons(Id ":c-li0-eq", Cons( Cons(l, _) , Cons( r , Nil))) when l = r ->
      Nil, true, ctx
  | Cons(Id ":c-int", Cons(Int _, Nil)) ->
      Nil, true, ctx
  | Cons(Id ":c-float", Cons(Flt _, Nil)) ->
      Nil, true, ctx
  | Cons(Id ":c-num", Cons(Int _, Nil))
  | Cons(Id ":c-num", Cons(Flt _, Nil)) ->
      Nil, true, ctx
  | Cons(Id ":c-str", Cons(Str _, Nil)) ->
      Nil, true, ctx
  | Cons(Id ":c-id", Cons(Id _, Nil)) ->
      Nil, true, ctx
  | Cons(Id ":c-colon-id", Cons(Id x, Nil)) when String.contains x ':' ->
      Nil, true, ctx
  | Cons(Id ":c-list", Cons(x, Nil)) when is_list x ->
      Nil, true, ctx
  | Cons(Id ":c-atm", Cons(Int _, Nil))
  | Cons(Id ":c-atm", Cons(Flt _, Nil))
  | Cons(Id ":c-atm", Cons(Str _, Nil))
  | Cons(Id ":c-atm", Cons(Id _, Nil)) ->
      Nil, true, ctx
  | Cons(Id i, r) when String.starts_with ~prefix:":" i ->
      let i = String.sub i 1 ((String.length i) - 1) in
      check_expands_macro i ctx r s
  | Cons(_) ->
      check_expands_macro "@" ctx s s
  | s -> li1 s, false, ctx
and check_expands_macro i ctx r s =
  let open Sexpr in
  let rec test opt ctx =
    begin match opt with
      [] -> li1 s, false, ctx
    | hd :: tl ->
        try
          let o = expand_macro r hd ctx, true, ctx in
          (* Format.printf "expanded: %s, with: %a\n%!" i Sexpr.pp_t r; *)
          o
        with Macro_Doesnt_App ->
          test tl ctx
    end
  in
  test (assoc_all i ctx) ctx
and match_macro_arg args margs ctx =
  let open Sexpr in
  match args, margs with
    Nil,        Nil                                  -> []
  | Cons(_),    Cons(Id(".."), Cons(Id(rem),Nil))    -> [rem, args]
  | Cons(v,rl), Cons(Id(k), rr)                      -> [k,li1 v] @ match_macro_arg rl rr ctx
  | Cons(v,rl), Cons(Cons(Id k, Cons(cst, Nil)), rr) ->
      let cst_exp = expand_macro_eval cst ["_", li1 v] in
      let cst_exp, _ = eval_while ctx cst_exp in
      if cst_exp <> Nil then
        raise Macro_Doesnt_App;
      [k, li1 v] @ match_macro_arg rl rr ctx
  | _                                                -> raise Macro_Doesnt_App
and macro_req_score args =
  let open Sexpr in
  match args with
    Nil                               -> 0
  | Cons(Id(".."), Cons(Id(rem),Nil)) -> 1
  | Cons(Id _, rr)                    -> 100 + macro_req_score rr
  | Cons(Cons(Id _, Cons(_,Nil)), rr) -> 110 + macro_req_score rr
  | _                                 -> raise Not_Valid_Macro_Arg_Syntax
and expand_macro_eval expr defs =
  let open Sexpr in
  let perfm expr =
    begin match expr with
      Id(i) when String.starts_with ~prefix:"'" i ->
        li1 @@ Id(String.sub i 1 (String.length i - 1))
    | Id(i) -> begin match List.assoc_opt i defs with
        Some x -> x
      | None   -> li1 expr
      end
    | x when is_list x -> li1 (expand_macro_eval x defs)
    | x -> li1 x
    end
  in
  if is_list expr then
    flat_map perfm expr
  else
    perfm expr
and expand_macro args macro ctx =
  let margs, mbody = macro in
  let defs = match_macro_arg args margs ctx in
  num_expands := !num_expands + 1;
  expand_macro_eval mbody defs
and eval_while ctx s = begin
  let s, ch, ctx = eval ctx s in
  if ch then
    Sexpr.aggr_flat_map eval_while ctx s
  else
    s, ctx
end
and eval_root_while ctx s = begin
  let s, ch, ctx = eval ctx s in
  if ch then begin
    (* Format.printf "%a\n%!" Sexpr.pp_t s; *)
    Sexpr.aggr_flat_map eval_root_while ctx s
  end else
    s, ctx
end

let do_eval s =
  eval_root_while [] s |> fst

let sparse src =
  Pc.doparse
    Format.err_formatter
    { source = src; path = "top" }
    (Sexpr.parse_top ())
  |> Option.get

let () = test "SexprMacroExp.do_eval.unchanged_num" @@ begin fun () ->
  let open Sexpr in
  let a = Int 1 in
  assert (do_eval a = li1 a)
end

let () = test "SexprMacroExp.do_eval.unchanged_list" ~depends:["SexprMacroExp.do_eval.unchanged_num"] @@ begin fun () ->
  let open Sexpr in
  let a = Cons(Int 1, Cons(Int 2, Nil)) in
  assert (do_eval a = li1 a)
end

let () = test "SexprMacroExp.do_eval.unchanged_nested_list" ~depends:["SexprMacroExp.do_eval.unchanged_list"] @@ begin fun () ->
  let open Sexpr in
  let a = Cons(Int 1, Cons( Cons(Int 10, Cons(Int 20, Nil)), Cons(Int 3, Nil))) in
  assert (do_eval a = li1 a)
end

let () = test "SexprMacroExp.do_eval.macro_no_body" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro one()) (:one)"
  and p = "()" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_body" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a() 1 2) (:a)"
  and p = "(1 2)" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_arg" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a(arg) 1 arg) (:a 2)"
  and p = "(1 2)" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_args" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a(arg arg2) arg2 1 arg) (:a 2 0)"
  and p = "(0 1 2)" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_varargs" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a(arg .. rem) arg 1 rem) (:a 0 2 3)"
  and p = "(0 1 2 3)" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_varargs_notmatch_none" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a(arg .. rem) arg 1 rem) (:a 0)"
  and p = "((:a 0))" in
  assert (do_eval (sparse a) = sparse p)
end

let () = test "SexprMacroExp.do_eval.macro_body_paren" ~depends:Sexpr.parser_tests @@ begin fun () ->
  let open Sexpr in
  let a = "(:macro a() (1)) (:a)"
  and p = "((1))" in
  assert (do_eval (sparse a) = sparse p)
end