open Ppxlib

let is_operator name =
  String.length name > 0 &&
  match name.[0] with
  | '!' | '$' | '%' | '&' | '*' | '+' | '-' | '.' | '/' | ':' | '<' | '=' | '>' | '?' | '@' | '^' | '|' | '~' -> true
  | _ -> false

let extract_fun_params expr =
  match expr.pexp_desc with
  | Pexp_function (params, None, Pfunction_body body) ->
      let pats = List.filter_map (fun p ->
        match p.pparam_desc with
        | Pparam_val (Nolabel, None, pat) -> Some pat
        | _ -> None
      ) params in
      (pats, body)
  | _ -> ([], expr)

let rec expr_to_ir ~loc (e : expression) : Ir.expr =
  match e.pexp_desc with
  | Pexp_ident { txt = Lident "pure"; _ } ->
      Location.raise_errorf ~loc "pure requires an argument"
  | Pexp_ident { txt = Lident "fail"; _ } ->
      Location.raise_errorf ~loc "fail requires an argument"
  | Pexp_ident { txt = Lident "any"; _ } ->
      Ir.Any { loc }
  | Pexp_ident { txt = Lident "eof"; _ } ->
      Ir.Eof { loc }
  | Pexp_ident { txt = Lident "cut"; _ } ->
      Ir.Cut { loc }
  | Pexp_ident { txt = Lident name; _ } ->
      Ir.Var { loc; name }

  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "pure"; _ }; _ }, [(Nolabel, arg)]) ->
      Ir.Pure { loc; value = arg }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "fail"; _ }; _ }, [(Nolabel, { pexp_desc = Pexp_constant (Pconst_string (msg, _, _)); _ })]) ->
      Ir.Fail { loc; msg }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "satisfy"; _ }; _ }, [(Nolabel, pred)]) ->
      Ir.Satisfy { loc; pred; label = None }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "token"; _ }; _ }, [(Nolabel, tok)]) ->
      Ir.Token { loc; tok }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "tokens"; _ }; _ }, [(Nolabel, toks)]) ->
      Ir.Tokens { loc; toks }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "one_of"; _ }; _ }, [(Nolabel, toks)]) ->
      Ir.OneOf { loc; toks }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "none_of"; _ }; _ }, [(Nolabel, toks)]) ->
      Ir.NoneOf { loc; toks }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "char"; _ }; _ }, [(Nolabel, c)]) ->
      Ir.Token { loc; tok = c }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "string"; _ }; _ }, [(Nolabel, s)]) ->
      Ir.Tokens { loc; toks = [%expr String.to_seq [%e s] |> List.of_seq] }

  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "attempt"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.Attempt { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "lookahead"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.Lookahead { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "not_followed_by"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.NotFollowedBy { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "many"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.Many { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "some"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.Some_ { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "optional"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.Optional { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "skip_many"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.SkipMany { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "skip_some"; _ }; _ }, [(Nolabel, p)]) ->
      Ir.SkipSome { loc; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "take_while"; _ }; _ }, [(Nolabel, pred)]) ->
      Ir.TakeWhile { loc; pred; at_least_one = false }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "take_while1"; _ }; _ }, [(Nolabel, pred)]) ->
      Ir.TakeWhile { loc; pred; at_least_one = true }

  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "option"; _ }; _ }, [(Nolabel, default); (Nolabel, p)]) ->
      Ir.Option { loc; default; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "sep_by"; _ }; _ }, [(Nolabel, p); (Nolabel, sep)]) ->
      Ir.SepBy { loc; p = expr_to_ir ~loc p; sep = expr_to_ir ~loc sep }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "sep_by1"; _ }; _ }, [(Nolabel, p); (Nolabel, sep)]) ->
      Ir.SepBy1 { loc; p = expr_to_ir ~loc p; sep = expr_to_ir ~loc sep }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "end_by"; _ }; _ }, [(Nolabel, p); (Nolabel, sep)]) ->
      Ir.EndBy { loc; p = expr_to_ir ~loc p; sep = expr_to_ir ~loc sep }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "end_by1"; _ }; _ }, [(Nolabel, p); (Nolabel, sep)]) ->
      Ir.EndBy1 { loc; p = expr_to_ir ~loc p; sep = expr_to_ir ~loc sep }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "many_till"; _ }; _ }, [(Nolabel, p); (Nolabel, e)]) ->
      Ir.ManyTill { loc; p = expr_to_ir ~loc p; end_ = expr_to_ir ~loc e }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "count"; _ }; _ }, [(Nolabel, n); (Nolabel, p)]) ->
      Ir.Count { loc; n; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "between"; _ }; _ }, [(Nolabel, o); (Nolabel, c); (Nolabel, p)]) ->
      Ir.Between { loc; open_ = expr_to_ir ~loc o; close = expr_to_ir ~loc c; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "chainl"; _ }; _ }, [(Nolabel, p); (Nolabel, op)]) ->
      Ir.ChainL { loc; p = expr_to_ir ~loc p; op = expr_to_ir ~loc op }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "chainl1"; _ }; _ }, [(Nolabel, p); (Nolabel, op)]) ->
      Ir.ChainL1 { loc; p = expr_to_ir ~loc p; op = expr_to_ir ~loc op }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "chainr"; _ }; _ }, [(Nolabel, p); (Nolabel, op)]) ->
      Ir.ChainR { loc; p = expr_to_ir ~loc p; op = expr_to_ir ~loc op }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "chainr1"; _ }; _ }, [(Nolabel, p); (Nolabel, op)]) ->
      Ir.ChainR1 { loc; p = expr_to_ir ~loc p; op = expr_to_ir ~loc op }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "choice"; _ }; _ }, [(Nolabel, list_expr)]) ->
      let ps = match list_expr.pexp_desc with
        | Pexp_construct ({ txt = Lident "[]"; _ }, None) -> []
        | Pexp_construct ({ txt = Lident "::"; _ }, _) ->
            let rec extract_list e = match e.pexp_desc with
              | Pexp_construct ({ txt = Lident "[]"; _ }, None) -> []
              | Pexp_construct ({ txt = Lident "::"; _ }, Some { pexp_desc = Pexp_tuple [h; t]; _ }) ->
                  expr_to_ir ~loc h :: extract_list t
              | _ -> Location.raise_errorf ~loc "choice expects a list"
            in
            extract_list list_expr
        | _ -> Location.raise_errorf ~loc "choice expects a list"
      in
      Ir.Choice { loc; ps; label = None }

  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident ">>="; _ }; _ }, [(Nolabel, p); (Nolabel, f)]) ->
      Ir.Bind { loc; p = expr_to_ir ~loc p; f }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "<|>"; _ }; _ }, [(Nolabel, p); (Nolabel, q)]) ->
      Ir.Alt { loc; p = expr_to_ir ~loc p; q = expr_to_ir ~loc q }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "<*>"; _ }; _ }, [(Nolabel, pf); (Nolabel, px)]) ->
      Ir.Apply { loc; pf = expr_to_ir ~loc pf; px = expr_to_ir ~loc px }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "*>"; _ }; _ }, [(Nolabel, p); (Nolabel, q)]) ->
      Ir.SeqRight { loc; p = expr_to_ir ~loc p; q = expr_to_ir ~loc q }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "<*"; _ }; _ }, [(Nolabel, p); (Nolabel, q)]) ->
      Ir.SeqLeft { loc; p = expr_to_ir ~loc p; q = expr_to_ir ~loc q }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "<$>"; _ }; _ }, [(Nolabel, f); (Nolabel, p)]) ->
      Ir.Map { loc; f; p = expr_to_ir ~loc p }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "<?>"; _ }; _ }, [(Nolabel, p); (Nolabel, { pexp_desc = Pexp_constant (Pconst_string (label, _, _)); _ })]) ->
      Ir.Label { loc; p = expr_to_ir ~loc p; label }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "memo"; _ }; _ }, [(Nolabel, { pexp_desc = Pexp_constant (Pconst_string (name, _, _)); _ }); (Nolabel, p)]) ->
      Ir.Memo { loc; name; p = expr_to_ir ~loc p }

  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident "fix"; _ }; _ }, [(Nolabel, fn)]) ->
      let (pats, body) = extract_fun_params fn in
      (match pats with
       | [pat] ->
           let names = extract_pattern_names pat in
           Ir.Fix { loc; arity = List.length names; names; body }
       | _ -> Location.raise_errorf ~loc "fix requires a single function argument")

  | Pexp_apply (_, _) when is_infix_chain e ->
      parse_infix_chain ~loc e

  | _ ->
      Location.raise_errorf ~loc "unsupported parser expression"

and is_infix_chain e =
  match e.pexp_desc with
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident op; _ }; _ }, [(Nolabel, _); (Nolabel, _)])
    when is_operator op -> true
  | _ -> false

and parse_infix_chain ~loc e =
  expr_to_ir ~loc e

and extract_pattern_names pat =
  match pat.ppat_desc with
  | Ppat_var { txt; _ } -> [txt]
  | Ppat_tuple pats -> List.concat_map extract_pattern_names pats
  | _ -> Location.raise_errorf ~loc:pat.ppat_loc "fix pattern must be variable or tuple of variables"

let expand_parser_expr ~ctxt expr =
  let loc = Expansion_context.Extension.extension_point_loc ctxt in
  let ir = expr_to_ir ~loc expr in
  let ir = Optimise.optimise ir in
  let errors = Check.check_expr [] ir in
  List.iter (fun e ->
    let err_loc = Check.error_loc e in
    let msg = Check.format_error e in
    Location.raise_errorf ~loc:err_loc "%s" msg
  ) errors;
  let body = Codegen.compile ~loc [] ir in
  [%expr fun (type inp) (module I : Combin.INPUT with type t = inp and type token = char) (input : inp) ->
    let _memo = Combin.Memo.create () in
    [%e body]]

let parser_extension =
  Extension.V3.declare
    "parser"
    Extension.Context.expression
    Ast_pattern.(single_expr_payload __)
    expand_parser_expr

let parser_rule = Context_free.Rule.extension parser_extension

let extract_let_bindings str_items =
  List.filter_map (fun item ->
    match item.pstr_desc with
    | Pstr_value (Nonrecursive, [vb]) ->
        (match vb.pvb_pat.ppat_desc with
         | Ppat_var { txt = name; _ } -> Some (name, vb.pvb_expr, vb.pvb_loc)
         | _ -> None)
    | _ -> None
  ) str_items

let expand_parser_stri ~ctxt str_items =
  let loc = Expansion_context.Extension.extension_point_loc ctxt in
  let bindings = extract_let_bindings str_items in
  let compiled = List.map (fun (name, expr, binding_loc) ->
    let ir = expr_to_ir ~loc:binding_loc expr in
    let ir = Optimise.optimise ir in
    let memoized_ir = Ir.Memo { loc = binding_loc; name; p = ir } in
    let errors = Check.check_expr [] memoized_ir in
    List.iter (fun e ->
      let err_loc = Check.error_loc e in
      let msg = Check.format_error e in
      Location.raise_errorf ~loc:err_loc "%s" msg
    ) errors;
    let body = Codegen.compile ~loc:binding_loc [] memoized_ir in
    let func = [%expr fun (type inp) (module I : Combin.INPUT with type t = inp and type token = char) _memo (input : inp) -> [%e body]] in
    Ast_builder.Default.value_binding ~loc:binding_loc
      ~pat:(Ast_builder.Default.pvar ~loc:binding_loc name)
      ~expr:func
  ) bindings in
  Ast_builder.Default.pstr_value ~loc Nonrecursive compiled

let parser_stri_extension =
  Extension.V3.declare
    "parser"
    Extension.Context.structure_item
    Ast_pattern.(pstr __)
    expand_parser_stri

let parser_stri_rule = Context_free.Rule.extension parser_stri_extension

let () = Driver.register_transformation ~rules:[parser_rule; parser_stri_rule] "ppx_combin"