open import "prelude.ml"
open import "./monads.ml"

private type consumed = C | E

private type error = Unexpected of string * list string

instance show error begin
  let show = function
    | Unexpected (s, []) -> s
    | Unexpected (s, ms) ->
      "Unexpected " ^ s ^ "\n"
        ^ "Expecting one of " ^ foldl (^) "" ms
end

private type result 'a =
  | Ok  of consumed * 'a    * string
  | Err of consumed * error * string

let join_consumed x y =
  match x with
  | C -> C
  | E -> y

type parser_t 'm 'a = private P of string -> 'm (result 'a)
type parser <- parser_t identity

instance functor 'm => functor (parser_t 'm) begin
  let f <$> P x = P @@ fun i ->
    flip map (x i) @@ fun p ->
    match p with
    | Ok (c, x, i) -> Ok (c, f x, i)
    | Err e -> Err e
end

instance monad 'm => applicative (parser_t 'm) begin
  let pure x = P (fun s -> pure (Ok (E, x, s)))
  let P f <*> P x = P @@ fun s ->
    let! f = f s
    match f with
    | Ok (c, f, s) ->
      let! x = x s
      match x with
      | Ok (c', x, s) -> pure @@ Ok (join_consumed c c', f x, s)
      | Err (c', p) -> pure @@ Err (join_consumed c c', p)
    | Err e -> pure @@ Err e
end

let x *> y = (fun _ x -> x) <$> x <*> y
let x <* y = (| const x y |)

instance monad 'm => monad (parser_t 'm) begin
  let P x >>= f = P @@ fun s ->
    let! x = x s
    match x with
    | Ok (c, x, s) ->
      let P kont = f x
      let! x = kont s
      match x with
      | Ok (c', x, s) -> pure @@ Ok (join_consumed c c', x, s)
      | Err (c', p) -> pure @@ Err (join_consumed c c', p)
    | Err e -> pure @@ Err e
end

let private fail e = P (fun s -> pure (Err (E, e, s)))

let empty : forall 'm 'a. monad 'm => parser_t 'm 'a =
  fail (Unexpected ("empty parse", []))

let unexpected e =
  fail (Unexpected (e, []))

let alt_err (Unexpected (u, xs)) (Unexpected (_, ys)) =
  Unexpected (u, xs ++ ys)

let P x <|> P y = P @@ fun s ->
  let! x = x s
  match x with
  | Ok x -> pure (Ok x)
  | Err (c, m, s) ->
    let! y = y s
    match y with
    | Ok (c', x, s) -> pure (Ok (join_consumed c c', x, s))
    | Err (c', m', s) -> pure (Err (join_consumed c c', alt_err m m', s))

let P x <+> y = P @@ fun s ->
  let! x = x s
  match x with
  | Ok x -> pure (Ok x)
  | Err (c, m, s) ->
    let P y = force y
    let! y = y s
    match y with
    | Ok (c', x, s) -> pure (Ok (join_consumed c c', x, s))
    | Err (c', m', s) -> pure (Err (join_consumed c c', alt_err m m', s))

private module S = import "lua/string.ml"

let char : forall 'm. applicative 'm => parser_t 'm string =
  P @@ fun s ->
    let x = S.substring s 1 1
    if x <> "" then
      let r = S.substring s 2 (S.length s)
      pure @@ Ok (C, x, r)
    else
      pure @@ Err (E, Unexpected ("end-of-file", ["character"]), s)

let eof : forall 'm. applicative 'm => parser_t 'm () =
  P @@ fun s ->
    if s == "" then
      pure @@ Ok (E, (), s)
    else
      pure @@ Err (E, Unexpected (S.substring s 1 1, ["end-of-file"]), s)

let satisfy p = P @@ fun s ->
  let x = S.substring s 1 1
  if x <> "" && p x then
    pure @@ Ok (C, x, S.substring s 2 (S.length s))
  else
    let m =
      if x == "" then
        "end of file"
      else x
    pure @@ Err (E, Unexpected (m, ["character"]), s)

let P k <?> m = P @@ fun s ->
  let! x = k s
  match x with
  | Ok e -> pure (Ok e)
  | Err (c, Unexpected (u, _), s) -> pure (Err (c, Unexpected (u, [m]), s))

let many_fold k z (P x) =
  let rec go consumed kont s =
    let! x = x s
    match x with
    | Err ((c, _, s) as p) ->
      match c with
      | C -> pure (Err p)
      | E -> kont consumed z s
    | Ok (c, x, s) ->
      match c with
      | E -> error @@ "many: got parser that accepts the empty string"
      | C -> go C (fun c -> kont c # k x) s
  P (go E (fun c z s -> pure (Ok (c, z, s))))

let many p = many_fold (::) [] p
let some p =
  let! x = p
  (x ::) <$> many p

let sep_by_1 sep p =
  let! x = p
  let! xs = many (sep *> p)
  pure (x :: xs)

let sep_by sep p = sep_by_1 sep p <|> pure []
external private val is_infix_of : string -> string -> bool =
  "function(s, s2) return s2:find(s) ~= nil end"

external private val is_prefix_of : string -> string -> bool =
  "function(s, s2) return (s2:find(s)) == 1 end"

let one_of chs =
  let len = S.length chs
  let rec loop ch i =
    if i > len then
      fail @@ Unexpected (ch, ["one of \"" ^ chs ^ "\""])
    else if ch == S.substring chs i i then
      pure ch
    else
      loop ch (i + 1)
  let! ch = char
  loop ch 1

let symbol str =
  let len = S.length str
  let rec loop acc i =
    if i > len then
      pure acc
    else
      let! c = char <?> S.substring str i i
      if c == S.substring str i i then
        loop (acc ^ c) (i + 1)
      else
        fail @@ Unexpected (acc ^ c, [S.substring str 1 i])
  loop "" 1

let between o c p =
  let! _ = o
  let! x = p
  let! _ = c
  pure x

let try (P x) = P @@ fun s ->
  let! x = x s
  match x with
  | Ok c -> pure (Ok c)
  | Err (_, p, _) -> pure (Err (E, p, s))

let optionally p = map Some (try p) <|> pure None


let rec sep_end_by_1 sep p =
  let! x = p
  ( let! _ = sep
    let! xs = sep_end_by sep p
    pure (x :: xs)
  ) <|> pure [x]
and sep_end_by sep p =
  sep_end_by_1 sep p <|> pure []


let parse (P x) s =
  let! x = x s
  match x with
  | Ok (_, x, r) -> pure @@ Right (x, r)
  | Err (_, m, _) -> pure @@ Left m

let run_parser p s =
  let Id x = parse p s
  x

let run_parser' (P x) s =
  let Id x = x s
  x

let lift m = P @@ fun s ->
  let! x = m
  pure @@ Ok (E, x, s)

let morph (k : forall 'a. 'm 'a -> 'n 'a) (P x) = P @@ fun s -> k (x s)
let void x = map (const ()) x