open Ir

type error =
  | Left_recursion of { loc : loc; name : string; path : string list }
  | Ambiguous_alt of { loc : loc }

let left_recursive_path env names expr =
  let rec nullable = function
    | Pure _ -> true
    | Fail _ -> false
    | Satisfy _ -> false
    | Any _ -> false
    | Eof _ -> true
    | Token _ -> false
    | Tokens _ -> false
    | OneOf _ -> false
    | NoneOf _ -> false
    | Bind { p; _ } -> nullable p
    | Map { p; _ } -> nullable p
    | Apply { pf; px; _ } -> nullable pf && nullable px
    | SeqLeft { p; q; _ } -> nullable p && nullable q
    | SeqRight { p; q; _ } -> nullable p && nullable q
    | Alt { p; q; _ } -> nullable p || nullable q
    | Attempt { p; _ } -> nullable p
    | Cut _ -> true
    | Lookahead _ -> true
    | NotFollowedBy _ -> true
    | Many _ -> true
    | Some_ { p; _ } -> nullable p
    | SkipMany _ -> true
    | SkipSome { p; _ } -> nullable p
    | TakeWhile { at_least_one; _ } -> not at_least_one
    | Optional _ -> true
    | Option _ -> true
    | SepBy _ -> true
    | SepBy1 { p; sep; _ } -> nullable p && nullable sep
    | EndBy _ -> true
    | EndBy1 { p; sep; _ } -> nullable p && nullable sep
    | ManyTill _ -> true
    | Count _ -> false
    | Between { open_; close; p; _ } -> nullable open_ && nullable p && nullable close
    | ChainL _ -> true
    | ChainL1 { p; _ } -> nullable p
    | ChainR _ -> true
    | ChainR1 { p; _ } -> nullable p
    | Choice { ps; _ } -> List.exists nullable ps
    | Label { p; _ } -> nullable p
    | Memo { p; _ } -> nullable p
    | Fix _ -> false
    | Var { name; _ } ->
        (match List.assoc_opt name env with
         | Some e -> nullable e
         | None -> false)
  in
  let rec check path = function
    | Var { name; _ } when List.mem name names ->
        Some (List.rev (name :: path))
    | Var { name; _ } ->
        (match List.assoc_opt name env with
         | Some e -> check (name :: path) e
         | None -> None)
    | Bind { p; _ } -> check path p
    | Map { p; _ } -> check path p
    | Apply { pf; px; _ } ->
        (match check path pf with
         | Some _ as r -> r
         | None -> if nullable pf then check path px else None)
    | SeqLeft { p; q; _ } ->
        (match check path p with
         | Some _ as r -> r
         | None -> if nullable p then check path q else None)
    | SeqRight { p; q; _ } ->
        (match check path p with
         | Some _ as r -> r
         | None -> if nullable p then check path q else None)
    | Alt { p; q; _ } ->
        (match check path p with
         | Some _ as r -> r
         | None -> check path q)
    | Attempt { p; _ } -> check path p
    | Lookahead { p; _ } -> check path p
    | NotFollowedBy { p; _ } -> check path p
    | Label { p; _ } -> check path p
    | Many { p; _ } -> check path p
    | Some_ { p; _ } -> check path p
    | SkipMany { p; _ } -> check path p
    | SkipSome { p; _ } -> check path p
    | TakeWhile _ -> None
    | Optional { p; _ } -> check path p
    | Option { p; _ } -> check path p
    | SepBy { p; _ } -> check path p
    | SepBy1 { p; _ } -> check path p
    | EndBy { p; _ } -> check path p
    | EndBy1 { p; _ } -> check path p
    | ManyTill { p; _ } -> check path p
    | Count { p; _ } -> check path p
    | Between { open_; _ } -> check path open_
    | ChainL { p; _ } -> check path p
    | ChainL1 { p; _ } -> check path p
    | ChainR { p; _ } -> check path p
    | ChainR1 { p; _ } -> check path p
    | Choice { ps; _ } -> List.find_map (check path) ps
    | Memo { p; _ } -> check path p
    | Fix _ -> None
    | Pure _ | Fail _ | Satisfy _ | Any _ | Eof _ | Token _ | Tokens _
    | OneOf _ | NoneOf _ | Cut _ -> None
  in
  check [] expr

let check_left_recursion env names expr =
  match left_recursive_path env names expr with
  | Some path -> Some (Left_recursion { loc = loc_of_expr expr; name = List.hd names; path })
  | None -> None

let is_wrapped_in_attempt = function
  | Attempt _ -> true
  | _ -> false

let rec check_expr env = function
  | Alt { loc; p; q } ->
      let p_first = First_set.compute env p in
      let q_first = First_set.compute env q in
      let errs =
        if is_wrapped_in_attempt p || is_wrapped_in_attempt q then
          []
        else if not (First_set.definitely_disjoint p_first q_first) then
          [Ambiguous_alt { loc }]
        else
          []
      in
      errs @ check_expr env p @ check_expr env q
  | Bind { p; _ } -> check_expr env p
  | Map { p; _ } -> check_expr env p
  | Apply { pf; px; _ } -> check_expr env pf @ check_expr env px
  | SeqLeft { p; q; _ } -> check_expr env p @ check_expr env q
  | SeqRight { p; q; _ } -> check_expr env p @ check_expr env q
  | Attempt { p; _ } -> check_expr env p
  | Lookahead { p; _ } -> check_expr env p
  | NotFollowedBy { p; _ } -> check_expr env p
  | Many { p; _ } -> check_expr env p
  | Some_ { p; _ } -> check_expr env p
  | SkipMany { p; _ } -> check_expr env p
  | SkipSome { p; _ } -> check_expr env p
  | TakeWhile _ -> []
  | Optional { p; _ } -> check_expr env p
  | Option { p; _ } -> check_expr env p
  | SepBy { p; sep; _ } -> check_expr env p @ check_expr env sep
  | SepBy1 { p; sep; _ } -> check_expr env p @ check_expr env sep
  | EndBy { p; sep; _ } -> check_expr env p @ check_expr env sep
  | EndBy1 { p; sep; _ } -> check_expr env p @ check_expr env sep
  | ManyTill { p; end_; _ } -> check_expr env p @ check_expr env end_
  | Count { p; _ } -> check_expr env p
  | Between { open_; close; p; _ } -> check_expr env open_ @ check_expr env close @ check_expr env p
  | ChainL { p; op; _ } -> check_expr env p @ check_expr env op
  | ChainL1 { p; op; _ } -> check_expr env p @ check_expr env op
  | ChainR { p; op; _ } -> check_expr env p @ check_expr env op
  | ChainR1 { p; op; _ } -> check_expr env p @ check_expr env op
  | Choice { ps; _ } -> List.concat_map (check_expr env) ps
  | Label { p; _ } -> check_expr env p
  | Memo { p; _ } -> check_expr env p
  | Pure _ | Fail _ | Satisfy _ | Any _ | Eof _ | Token _ | Tokens _
  | OneOf _ | NoneOf _ | Cut _ | Fix _ | Var _ -> []

let error_loc = function
  | Left_recursion { loc; _ } -> loc
  | Ambiguous_alt { loc } -> loc

let format_error = function
  | Left_recursion { name; path; _ } ->
      Printf.sprintf "left recursion detected: %s via %s"
        name (String.concat " -> " path)
  | Ambiguous_alt _ ->
      "ambiguous alternation: branches may have overlapping first sets, wrap one in 'attempt'"