open Ppxlib
open Ast_builder.Default
open Ir

let fresh =
  let counter = ref 0 in
  fun prefix ->
    incr counter;
    Printf.sprintf "%s_%d" prefix !counter

let error_expr ~loc expected =
  [%expr Error { Combin.span = { Combin.start_pos = I.position input; end_pos = I.position input }; expected = [%e expected] }]

let error_expr_span ~loc start_var expected =
  [%expr Error { Combin.span = { Combin.start_pos = [%e start_var]; end_pos = I.position input }; expected = [%e expected] }]

let expected_list ~loc strs =
  elist ~loc (List.map (estring ~loc) strs)

let ok_expr ~loc value input_e =
  [%expr Ok ([%e value], [%e input_e])]

let merge_errors_expr ~loc e1 e2 =
  [%expr Combin.merge_errors [%e e1] [%e e2]]

let pvar_s ~loc s = ppat_var ~loc { txt = s; loc }
let evar_s ~loc s = evar ~loc s

let mk_let_rec ~loc name params body rest =
  let param_pats = List.map (pvar_s ~loc) params in
  let func = List.fold_right (fun p e -> pexp_fun ~loc Nolabel None p e) param_pats body in
  let vb = value_binding ~loc ~pat:(pvar_s ~loc name) ~expr:func in
  pexp_let ~loc Recursive [vb] rest

let rec compile ~loc env expr =
  match expr with
  | Pure { value; _ } ->
      [%expr Ok ([%e value], input)]

  | Fail { msg; _ } ->
      error_expr ~loc (elist ~loc [estring ~loc msg])

  | Satisfy { pred; label; loc } ->
      let expected = match label with
        | Some l -> elist ~loc [estring ~loc l]
        | None -> [%expr ["<satisfy>"]]
      in
      [%expr
        match I.peek input with
        | Some c when [%e pred] c -> Ok (c, I.advance input)
        | _ -> [%e error_expr ~loc expected]]

  | Any { loc } ->
      [%expr
        match I.peek input with
        | Some c -> Ok (c, I.advance input)
        | None -> [%e error_expr ~loc (elist ~loc [estring ~loc "<any>"])]]

  | Eof { loc } ->
      [%expr
        match I.peek input with
        | None -> Ok ((), input)
        | Some _ -> [%e error_expr ~loc (elist ~loc [estring ~loc "<eof>"])]]

  | Token { tok; loc } ->
      [%expr
        match I.peek input with
        | Some t when t = [%e tok] -> Ok (t, I.advance input)
        | _ -> [%e error_expr ~loc [%expr [I.show_token [%e tok]]]]]

  | Tokens { toks; loc } ->
      [%expr
        let rec loop toks_list inp =
          match toks_list with
          | [] -> Ok ([%e toks], inp)
          | t :: rest ->
              match I.peek inp with
              | Some t' when t = t' -> loop rest (I.advance inp)
              | _ -> [%e error_expr ~loc [%expr [String.concat "" (List.map I.show_token [%e toks])]]]
        in
        loop [%e toks] input]

  | OneOf { toks; loc } ->
      [%expr
        match I.peek input with
        | Some t when List.mem t [%e toks] -> Ok (t, I.advance input)
        | _ -> [%e error_expr ~loc [%expr List.map I.show_token [%e toks]]]]

  | NoneOf { toks; loc } ->
      [%expr
        match I.peek input with
        | Some t when not (List.mem t [%e toks]) -> Ok (t, I.advance input)
        | _ -> [%e error_expr ~loc (elist ~loc [estring ~loc "<none_of>"])]]

  | Bind { p; f; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) ->
            let input = inp in
            [%e f] x input
        | Error e -> Error e]

  | Map { f; p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) -> Ok ([%e f] x, inp)
        | Error e -> Error e]

  | Apply { pf; px; loc } ->
      let pf_code = compile ~loc env pf in
      let px_code = compile ~loc env px in
      [%expr
        match [%e pf_code] with
        | Ok (f, inp1) ->
            let input = inp1 in
            (match [%e px_code] with
             | Ok (x, inp2) -> Ok (f x, inp2)
             | Error e -> Error e)
        | Error e -> Error e]

  | SeqLeft { p; q; loc } ->
      let p_code = compile ~loc env p in
      let q_code = compile ~loc env q in
      [%expr
        match [%e p_code] with
        | Ok (x, inp1) ->
            let input = inp1 in
            (match [%e q_code] with
             | Ok (_, inp2) -> Ok (x, inp2)
             | Error e -> Error e)
        | Error e -> Error e]

  | SeqRight { p; q; loc } ->
      let p_code = compile ~loc env p in
      let q_code = compile ~loc env q in
      [%expr
        match [%e p_code] with
        | Ok (_, inp1) ->
            let input = inp1 in
            [%e q_code]
        | Error e -> Error e]

  | Alt { p; q; loc } ->
      let p_first = First_set.compute [] p in
      let q_first = First_set.compute [] q in
      if First_set.definitely_disjoint p_first q_first &&
         First_set.all_static p_first && First_set.all_static q_first then
        compile_committed_alt ~loc env p q p_first q_first
      else
        compile_backtrack_alt ~loc env p q

  | Attempt { p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        let saved = input in
        match [%e p_code] with
        | Ok _ as r -> r
        | Error e -> Error { e with Combin.span = { e.Combin.span with Combin.start_pos = I.position saved } }]

  | Cut { loc } ->
      [%expr Ok ((), input)]

  | Lookahead { p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        let saved = input in
        match [%e p_code] with
        | Ok (v, _) -> Ok (v, saved)
        | Error e -> Error e]

  | NotFollowedBy { p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        let saved = input in
        match [%e p_code] with
        | Ok _ -> [%e error_expr ~loc (elist ~loc [estring ~loc "<not_followed_by>"])]
        | Error _ -> Ok ((), saved)]

  | Many { p; loc } ->
      (match p with
       | Token { tok; _ } ->
           [%expr
             let rec loop acc input =
               match I.peek input with
               | Some t when t = [%e tok] -> loop (t :: acc) (I.advance input)
               | _ -> Ok (List.rev acc, input)
             in
             loop [] input]
       | Satisfy { pred; _ } ->
           [%expr
             let rec loop acc input =
               match I.peek input with
               | Some t when [%e pred] t -> loop (t :: acc) (I.advance input)
               | _ -> Ok (List.rev acc, input)
             in
             loop [] input]
       | Any _ ->
           [%expr
             let rec loop acc input =
               match I.peek input with
               | Some t -> loop (t :: acc) (I.advance input)
               | None -> Ok (List.rev acc, input)
             in
             loop [] input]
       | _ ->
           let p_code = compile ~loc env p in
           [%expr
             let rec loop acc input =
               match [%e p_code] with
               | Ok (x, inp) -> loop (x :: acc) inp
               | Error _ -> Ok (List.rev acc, input)
             in
             loop [] input])

  | Some_ { p; loc } ->
      (match p with
       | Token { tok; _ } ->
           [%expr
             match I.peek input with
             | Some t when t = [%e tok] ->
                 let rec loop acc input =
                   match I.peek input with
                   | Some t when t = [%e tok] -> loop (t :: acc) (I.advance input)
                   | _ -> Ok (List.rev acc, input)
                 in
                 loop [t] (I.advance input)
             | _ -> [%e error_expr ~loc [%expr [I.show_token [%e tok]]]]]
       | Satisfy { pred; label; _ } ->
           let expected = match label with
             | Some l -> elist ~loc [estring ~loc l]
             | None -> [%expr ["<satisfy>"]]
           in
           [%expr
             match I.peek input with
             | Some t when [%e pred] t ->
                 let rec loop acc input =
                   match I.peek input with
                   | Some t when [%e pred] t -> loop (t :: acc) (I.advance input)
                   | _ -> Ok (List.rev acc, input)
                 in
                 loop [t] (I.advance input)
             | _ -> [%e error_expr ~loc expected]]
       | Any _ ->
           [%expr
             match I.peek input with
             | Some t ->
                 let rec loop acc input =
                   match I.peek input with
                   | Some t -> loop (t :: acc) (I.advance input)
                   | None -> Ok (List.rev acc, input)
                 in
                 loop [t] (I.advance input)
             | None -> [%e error_expr ~loc (elist ~loc [estring ~loc "<any>"])]]
       | _ ->
           let p_code = compile ~loc env p in
           [%expr
             match [%e p_code] with
             | Ok (x, inp) ->
                 let rec loop acc input =
                   match [%e p_code] with
                   | Ok (y, inp2) -> loop (y :: acc) inp2
                   | Error _ -> Ok (List.rev acc, input)
                 in
                 loop [x] inp
             | Error e -> Error e])

  | SkipMany { p; loc } ->
      (match p with
       | Token { tok; _ } ->
           [%expr
             let rec loop input =
               match I.peek input with
               | Some t when t = [%e tok] -> loop (I.advance input)
               | _ -> Ok ((), input)
             in
             loop input]
       | Satisfy { pred; _ } ->
           [%expr
             let rec loop input =
               match I.peek input with
               | Some t when [%e pred] t -> loop (I.advance input)
               | _ -> Ok ((), input)
             in
             loop input]
       | Any _ ->
           [%expr
             let rec loop input =
               match I.peek input with
               | Some _ -> loop (I.advance input)
               | None -> Ok ((), input)
             in
             loop input]
       | _ ->
           let p_code = compile ~loc env p in
           [%expr
             let rec loop input =
               match [%e p_code] with
               | Ok (_, inp) -> loop inp
               | Error _ -> Ok ((), input)
             in
             loop input])

  | SkipSome { p; loc } ->
      (match p with
       | Token { tok; _ } ->
           [%expr
             match I.peek input with
             | Some t when t = [%e tok] ->
                 let rec loop input =
                   match I.peek input with
                   | Some t when t = [%e tok] -> loop (I.advance input)
                   | _ -> Ok ((), input)
                 in
                 loop (I.advance input)
             | _ -> [%e error_expr ~loc [%expr [I.show_token [%e tok]]]]]
       | Satisfy { pred; label; _ } ->
           let expected = match label with
             | Some l -> elist ~loc [estring ~loc l]
             | None -> [%expr ["<satisfy>"]]
           in
           [%expr
             match I.peek input with
             | Some t when [%e pred] t ->
                 let rec loop input =
                   match I.peek input with
                   | Some t when [%e pred] t -> loop (I.advance input)
                   | _ -> Ok ((), input)
                 in
                 loop (I.advance input)
             | _ -> [%e error_expr ~loc expected]]
       | _ ->
           let p_code = compile ~loc env p in
           [%expr
             match [%e p_code] with
             | Ok (_, inp) ->
                 let rec loop input =
                   match [%e p_code] with
                   | Ok (_, inp2) -> loop inp2
                   | Error _ -> Ok ((), input)
                 in
                 loop inp
             | Error e -> Error e])

  | TakeWhile { pred; at_least_one; loc } ->
      if at_least_one then
        [%expr
          match I.peek input with
          | Some t when [%e pred] t ->
              let buf = Buffer.create 16 in
              let rec loop input =
                match I.peek input with
                | Some t when [%e pred] t ->
                    Buffer.add_char buf t;
                    loop (I.advance input)
                | _ -> Ok (Buffer.contents buf, input)
              in
              Buffer.add_char buf t;
              loop (I.advance input)
          | _ -> [%e error_expr ~loc [%expr ["<take_while1>"]]]]
      else
        [%expr
          let buf = Buffer.create 16 in
          let rec loop input =
            match I.peek input with
            | Some t when [%e pred] t ->
                Buffer.add_char buf t;
                loop (I.advance input)
            | _ -> Ok (Buffer.contents buf, input)
          in
          loop input]

  | Optional { p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) -> Ok (Some x, inp)
        | Error _ -> Ok (None, input)]

  | Option { default; p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) -> Ok (x, inp)
        | Error _ -> Ok ([%e default], input)]

  | SepBy { p; sep; loc } ->
      let sepby1 = compile ~loc env (SepBy1 { p; sep; loc }) in
      [%expr
        match [%e sepby1] with
        | Ok _ as r -> r
        | Error _ -> Ok ([], input)]

  | SepBy1 { p; sep; loc } ->
      let p_code = compile ~loc env p in
      let sep_code = compile ~loc env sep in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) ->
            let rec loop acc input =
              match [%e sep_code] with
              | Ok (_, inp2) ->
                  let input = inp2 in
                  (match [%e p_code] with
                   | Ok (y, inp3) -> loop (y :: acc) inp3
                   | Error e -> Error e)
              | Error _ -> Ok (List.rev acc, input)
            in
            loop [x] inp
        | Error e -> Error e]

  | EndBy { p; sep; loc } ->
      let endby1 = compile ~loc env (EndBy1 { p; sep; loc }) in
      [%expr
        match [%e endby1] with
        | Ok _ as r -> r
        | Error _ -> Ok ([], input)]

  | EndBy1 { p; sep; loc } ->
      let p_code = compile ~loc env p in
      let sep_code = compile ~loc env sep in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) ->
            let input = inp in
            (match [%e sep_code] with
             | Ok (_, inp2) ->
                 let rec loop acc input =
                   match [%e p_code] with
                   | Ok (y, inp3) ->
                       let input = inp3 in
                       (match [%e sep_code] with
                        | Ok (_, inp4) -> loop (y :: acc) inp4
                        | Error e -> Error e)
                   | Error _ -> Ok (List.rev acc, input)
                 in
                 loop [x] inp2
             | Error e -> Error e)
        | Error e -> Error e]

  | ManyTill { p; end_; loc } ->
      let p_code = compile ~loc env p in
      let end_code = compile ~loc env end_ in
      [%expr
        let rec loop acc input =
          match [%e end_code] with
          | Ok (_, inp) -> Ok (List.rev acc, inp)
          | Error _ ->
              match [%e p_code] with
              | Ok (x, inp) -> loop (x :: acc) inp
              | Error e -> Error e
        in
        loop [] input]

  | Count { n; p; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        let rec loop n acc input =
          if n <= 0 then Ok (List.rev acc, input)
          else
            match [%e p_code] with
            | Ok (x, inp) -> loop (n - 1) (x :: acc) inp
            | Error e -> Error e
        in
        loop [%e n] [] input]

  | Between { open_; close; p; loc } ->
      compile ~loc env (SeqRight { loc; p = open_; q = SeqLeft { loc; p; q = close } })

  | ChainL { p; op; loc } ->
      let chainl1 = compile ~loc env (ChainL1 { p; op; loc }) in
      [%expr
        match [%e chainl1] with
        | Ok _ as r -> r
        | Error _ -> Ok ([], input)]

  | ChainL1 { p; op; loc } ->
      let p_code = compile ~loc env p in
      let op_code = compile ~loc env op in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) ->
            let rec loop acc input =
              match [%e op_code] with
              | Ok (f, inp2) ->
                  let input = inp2 in
                  (match [%e p_code] with
                   | Ok (y, inp3) -> loop (f acc y) inp3
                   | Error e -> Error e)
              | Error _ -> Ok (acc, input)
            in
            loop x inp
        | Error e -> Error e]

  | ChainR { p; op; loc } ->
      let chainr1 = compile ~loc env (ChainR1 { p; op; loc }) in
      [%expr
        match [%e chainr1] with
        | Ok _ as r -> r
        | Error _ -> Ok ([], input)]

  | ChainR1 { p; op; loc } ->
      let p_code = compile ~loc env p in
      let op_code = compile ~loc env op in
      [%expr
        match [%e p_code] with
        | Ok (x, inp) ->
            let rec loop stack input =
              match [%e op_code] with
              | Ok (f, inp2) ->
                  let input = inp2 in
                  (match [%e p_code] with
                   | Ok (y, inp3) -> loop ((f, y) :: stack) inp3
                   | Error e -> Error e)
              | Error _ ->
                  let rec fold v = function
                    | [] -> v
                    | (f, y) :: rest -> fold (f v y) rest
                  in
                  Ok (fold x (List.rev stack), input)
            in
            loop [] inp
        | Error e -> Error e]

  | Choice { ps; label; loc } ->
      let rec build_choice = function
        | [] -> error_expr ~loc (match label with Some l -> elist ~loc [estring ~loc l] | None -> elist ~loc [estring ~loc "<choice>"])
        | [p] -> compile ~loc env p
        | p :: rest ->
            let p_code = compile ~loc env p in
            let rest_code = build_choice rest in
            [%expr
              let saved = input in
              match [%e p_code] with
              | Ok _ as r -> r
              | Error e1 ->
                  let input = saved in
                  match [%e rest_code] with
                  | Ok _ as r -> r
                  | Error e2 -> Error (Combin.merge_errors e1 e2)]
      in
      build_choice ps

  | Label { p; label; loc } ->
      let p_code = compile ~loc env p in
      [%expr
        match [%e p_code] with
        | Ok _ as r -> r
        | Error e -> Error { e with Combin.expected = [[%e estring ~loc label]] }]

  | Memo { name; p; loc } ->
      let p_code = compile ~loc env p in
      let name_expr = estring ~loc name in
      [%expr
        let pos = I.position input in
        match Combin.Memo.find _memo [%e name_expr] pos with
        | Some r -> r
        | None ->
            let r = [%e p_code] in
            Combin.Memo.add _memo [%e name_expr] pos r;
            r]

  | Fix { body; _ } ->
      [%expr
        let rec parser_fix = [%e body] in
        parser_fix input]

  | Var { name; loc } ->
      [%expr [%e evar ~loc name] (module I) _memo input]

and compile_committed_alt ~loc env p q p_first q_first =
  let p_code = compile ~loc env p in
  let q_code = compile ~loc env q in
  let build_guard first =
    match first.First_set.firsts with
    | [First_set.Token tok] ->
        [%expr match I.peek input with Some t when t = [%e tok] -> true | _ -> false]
    | [First_set.OneOf toks] ->
        [%expr match I.peek input with Some t when List.mem t [%e toks] -> true | _ -> false]
    | _ -> [%expr true]
  in
  let p_guard = build_guard p_first in
  let q_guard = build_guard q_first in
  [%expr
    if [%e p_guard] then [%e p_code]
    else if [%e q_guard] then [%e q_code]
    else [%e error_expr ~loc (elist ~loc [estring ~loc "<alt>"])]]

and compile_backtrack_alt ~loc env p q =
  let p_code = compile ~loc env p in
  let q_code = compile ~loc env q in
  [%expr
    let saved = input in
    match [%e p_code] with
    | Ok _ as r -> r
    | Error e1 ->
        let input = saved in
        match [%e q_code] with
        | Ok _ as r -> r
        | Error e2 -> Error (Combin.merge_errors e1 e2)]

let compile_def env { name; expr; loc } =
  let body = compile ~loc env expr in
  let func = [%expr fun (type inp) (module I : Combin.INPUT with type t = inp and type token = char) input -> [%e body]] in
  value_binding ~loc ~pat:(pvar ~loc name) ~expr:func

let compile_mutual_def env { names; expr; loc } =
  let body = compile ~loc env expr in
  let func = [%expr fun (type inp) (module I : Combin.INPUT with type t = inp and type token = char) input -> [%e body]] in
  let pat = ppat_tuple ~loc (List.map (pvar ~loc) names) in
  value_binding ~loc ~pat ~expr:func