let test = Testing.test;;

module SExpr = struct
  type t =
    Nil
  | Int  of int
  | Flt  of float
  | Cons of t * t
  | Id   of string
  | Str  of string

  let li1 s =
    Cons(s, Nil)

  exception Not_SExpr_List

  let rec append tl hd =
    match hd with
      Nil          -> tl
    | Cons(x, rem) -> Cons(x, append tl rem)
    | _            -> raise Not_SExpr_List

  let () = test "SExpr.append" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Cons(Int 3, Nil)))
    and b = Cons(Int 4, Cons(Int 5, Nil))
    and ab = Cons(Int 1, Cons(Int 2, Cons(Int 3, Cons(Int 4, Cons(Int 5, Nil))))) in
    assert (append b a = ab)
  end

  let rec rem_nil x =
    match x with
      Cons(x, Nil) -> rem_nil x
    | Cons(l, r)   -> Cons(l, rem_nil r)
    | x            -> x

  let () = test "SExpr.rem_nil" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Cons(Int 3, Nil)))
    and p = Cons(Int 1, Cons(Int 2, Int 3)) in
    assert (rem_nil a = p)
  end

  let rec flat_map fn li =
    match li with
      Nil          -> Nil
    | Cons(x, rem) -> append (flat_map fn rem) (fn x)
    | _            -> raise Not_SExpr_List

  let () = test "SExpr.flat_map.0" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Cons(Int 3, Nil)))
    and p = Nil in
    assert (flat_map (fun _ -> Nil) a = p)
  end

  let () = test "SExpr.flat_map.1" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Cons(Int 3, Nil)))
    and p = Cons(Int 2, Cons(Int 3, Cons(Int 4, Nil))) in
    assert (flat_map (fun x -> Cons((match x with Int(n) -> Int(n + 1) | x -> x) , Nil)) a = p)
  end

  let () = test "SExpr.flat_map.2" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Nil))
    and p = Cons(Int 1, Cons(Int 2, Cons(Int 2, Cons(Int 3, Nil)))) in
    assert (flat_map (fun x -> Cons(x, Cons((match x with Int(n) -> Int(n + 1) | x -> x) , Nil))) a = p)
  end

  let rec aggr_flat_map fn acc li =
    match li with
      Nil          -> Nil, acc
    | Cons(x, rem) ->
        let x, acc = fn acc x in
        let vv, acc = aggr_flat_map fn acc rem in
        append vv x, acc
    | _  -> raise Not_SExpr_List

  let () = test "SExpr.aggr_flat_map.0" @@ begin fun () ->
    let a = Cons(Int 1, Cons(Int 2, Nil))
    and p = Cons(Int 1, Cons(Int 2, Cons(Int 2, Cons(Int 4, Nil)))) in
    assert (fst (aggr_flat_map (fun aggr x -> Cons(x, Cons((match x with Int(n) -> Int(n + aggr) | x -> x) , Nil)) , aggr + 1) 1 a) = p)
  end

  (* a "list" here is: 
   *   ( e1 . ( e2 . ( e3 . NIL ) ) )
   * or even just:
   *   NIL
   *)
  let rec is_list s =
    match  s with
      Nil       -> true
    | Cons(a,b) -> is_list b
    | _         -> false

  let needs_surround_pipes x =
       String.contains x ' '
    || String.contains x '\n'
    || String.contains x '\r'
    || String.contains x '\t'
    || String.contains x '\b'

  let rec pp_t out (s:t) =
    let open Format in
    match s with
      Flt x -> fprintf out "%f" x
    | Int x -> fprintf out "%d" x
    | Str x ->
        fprintf out "\"";
        x |> String.iter begin fun x ->
          match x with
            '\\' -> fprintf out "\\\\"
          | '\r' -> fprintf out "\\r"
          | '\b' -> fprintf out "\\b"
          | '\n' -> fprintf out "\\n"
          | '\t' -> fprintf out "\\t"
          |  _   -> fprintf out "%c"  x
        end;
        fprintf out "\""
    | Id  x ->
        if needs_surround_pipes x then
          fprintf out "|%s|" x
        else
          fprintf out "%s" x
    | Nil -> fprintf out "()"
    | Cons(a,b) ->
        if is_list b then begin
          fprintf out "@[(@[";
          pp_closed out s
        end else begin
          fprintf out "@[(@[";
          pp_t out a;
          fprintf out "@ . ";
          pp_t out b;
          fprintf out "@[)@["
        end
  and pp_closed out (s:t) =
    let open Format in
    match s with
      Flt x -> fprintf out "%f" x
    | Int x -> fprintf out "%d" x
    | Str x -> fprintf out "\"%s\"" x
    | Id  x ->
        if needs_surround_pipes x then
          fprintf out "|%s|" x
        else
          fprintf out "%s" x
    | Nil -> fprintf out "@])@]"
    | Cons(a,b) ->
        pp_t out a;
        (match b with
          Nil -> ()
        | _   -> fprintf out "@ ");
        pp_closed out b

  let parse_inline () : Pc.cursor -> t * Pc.parse_result =
    let open Pc in
    let comment = chain (just ";") (repeat (also any (inv (just "\n")))) in
    let pad = either (ignore single_white) (ignore comment) in
    let pad = repeat pad in
    let padded p = (then_ignore (ignore_then pad p) pad) in
    recursive (fun expr ->
      let x_cons = nd_of3
           (just "(")
           begin
               chain expr (
                  ignore_then
                    (just ".")
                    expr
                  |> many)
           |> map (fun (head,tail) ->
               let li = List.rev (head :: tail) in
               let init = List.hd li in
               let others = List.rev(List.tl li) in
               List.fold_right (fun x acc -> Cons(x,acc)) others init)
           end
           (just ")")
      in
      let x_list = (nd_of3
        (padded (just "("))
        (repeat expr
        |> map (fun li ->
          List.fold_right (fun x acc -> Cons(x,acc)) li Nil))
        (just ")"))
      in
      let id_char = also any (inv begin
        (*-->> *) (ignore (just ";"))
        |> either (ignore single_white)
        |> either (ignore (just "("))
        |> either (ignore (just ")"))
        |> either (ignore (just "NIL"))
        |> either (ignore (just "|"))
        |> either (ignore (just "\""))
      end) in
      let escape = begin
        (*-->> *) also any (just "\"" |> inv)
        |> either (just "\\\\" |> set '\\')
        |> either (just "\\r"  |> set '\r')
        |> either (just "\\t"  |> set '\t')
        |> either (just "\\b"  |> set '\b')
        |> either (just "\\n"  |> set '\n')
      end in
          (* === expr === *)
      (*-->> *) (just "NIL" |> set Nil)
      |> either x_cons
      |> either x_list
      |> either (id_char |> many |> map
           (fun s -> Id(s |> List.to_seq |> String.of_seq )))
      |> either (nd_of3
           (just "|")
           begin also any (inv (just "|"))
              |> repeat
              |> map (fun s ->
                Id(s |> List.to_seq |> String.of_seq ))
           end
           (just "|")
         )
      |> either (p_int |> map (fun v -> Int(v)))
      |> either (p_flt |> map (fun v -> Flt(v)))
      |> either (nd_of3
           (just "\"")
           (escape
            |> repeat
            |> map (fun v -> Str(v |> List.to_seq |> String.of_seq)))
           (just "\""))
      |> padded
    )

  let parse_inline_end () =
    let open Pc in
    then_ignore (parse_inline ()) ex_end

  let parse_top () =
    let open Pc in
    parse_inline ()
      |> until ex_end
      |> map (fun li ->
          List.fold_right (fun x acc -> Cons(x,acc)) li Nil)

  let parser_tests = []
end


module SExprMacroExp = struct
  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
end


let read_all_stdin () =
  let buf = Buffer.create 4096 in
  try
    while true do
      let line = input_line stdin in
      Buffer.add_string buf line;
      Buffer.add_char buf '\n';
    done;
    Buffer.contents buf
  with End_of_file ->
    Buffer.contents buf


let () =
  let src = read_all_stdin () in
  let v = SExprMacroExp.sparse src in
  let v = SExprMacroExp.do_eval v in

  Format.printf "; number of macro expansions: %d\n%!" !SExprMacroExp.num_expands;

  Format.set_margin 40;
  Format.printf "%a@.@?" SExpr.pp_t v;

  exit 0