open Lexer

type parser_context = {
  seq: Lexer.token Seq.t;
  errors: string list;
}

(* The parser is a function that takes a parser_context and returns an option of a tuple of a value and a parser_context.*)
type 'a parser = parser_context -> ('a * parser_context) option

let return (a: 'a) = fun (ctx: parser_context) -> Some (a, ctx)
let stop = fun (_: parser_context) -> None

let fmap (f: 'a -> 'b) (p: 'a parser): 'b parser = fun (ctx: parser_context) ->
  match p ctx with
  | Some (a, ctx') -> Some (f a, ctx')
  | None -> None

let bind (a: 'a parser) (b:'a -> 'b parser) = fun (ctx: parser_context) ->
  let p = a ctx in
  match p with
  | Some (a', ctx') -> b a' ctx'
  | None -> None

let (>>=) = bind
let (let*) = bind

let or_parser (a: 'a parser) (b: 'a parser): 'a parser = fun (ctx: parser_context) ->
  match a ctx with
  | Some _ as res -> res
  | None -> b ctx

let (<|>) = or_parser

let peek_token: token parser = fun (ctx: parser_context) -> 
  Seq.uncons ctx.seq |> Option.map (fun (t,_) -> (t,ctx))

let next_token: token parser = fun (ctx: parser_context) -> 
  Seq.uncons ctx.seq |> Option.map (fun (t, s) -> (t, 
  { ctx with seq = s}
))

let match_token  (tt: token_type) : token parser = 
  let* t = next_token in
  if t.token_type = tt then
    return t
  else
    stop 

let zero_or_one (p: 'a parser): ('a option) parser = fun (ctx) ->
  match p ctx with
    | Some (a, ctx') -> Some (Some a, ctx')
    | None -> Some (None, ctx)

let rec many (p: 'a parser): 'a list parser =
  let* a = zero_or_one p in
  match a with
  | Some a' -> (
    let* as' = many p in
    return (a'::as')
  )
  | None -> return []

let many1 (p: 'a parser): 'a list parser =
  let* a = p in
  let* as' = many p in
  return (a::as')

(* 
BNF:
  let_expr ::= let identifier = expr in expr
  fun_expr ::= fun identifier -> expr
  if_expr ::= if expr then expr else expr
  factor ::= (expr) | identifier | number
  call_expr ::= factor | factor factor
  level1 ::= call_expr | level1 + call_expr | level1 - call_expr
  level2 ::= level2 * level1 | level2 / level1 | level2 % level1 | level1  
  level3 ::= level2 ^ level3 | level2
  expr ::= let_expr | fun_expr | if_expr | level3
*)

type let_expr_tree = Let of string * expr_tree * expr_tree
and fun_expr_tree = Fun of string * expr_tree
and if_expr_tree = If of expr_tree * expr_tree * expr_tree
and call_expr_tree = Call of expr_tree * expr_tree
and expr_tree = 
  | LetExpr of let_expr_tree
  | FunExpr of fun_expr_tree
  | IfExpr of if_expr_tree
  | CallExpr of call_expr_tree
  | BinOpExpr of Lexer.op_type * expr_tree * expr_tree
  | MonoOpExpr of Lexer.op_type * expr_tree
  | Identifier of string
  | Number of int

let expr2str (e: expr_tree): string = 
  let tab n = String.make (n * 2) ' ' in
  let rec aux e depth = 
    match e with
    | LetExpr (Let (id, e1, e2)) -> Printf.sprintf "let %s = %s in\n%s%s" id (aux e1 depth) (tab depth) (aux e2 (depth+1))
    | FunExpr (Fun (id, e)) -> Printf.sprintf "fun %s ->\n%s%s" id (tab depth) (aux e (depth+1))
    | IfExpr (If (e1, e2, e3)) -> Printf.sprintf "if %s then\n%s%selse\n%s%s" (aux e1 depth) (tab depth) (aux e2 depth) (tab depth) (aux e3 depth)
    | CallExpr (Call (e1, e2)) -> Printf.sprintf "%s %s" (aux e1 depth) (aux e2 depth)
    | BinOpExpr (op, e1, e2) -> Printf.sprintf "%s %s %s" (aux e1 depth) (Lexer.op2str op) (aux e2 depth)
    | MonoOpExpr (op, e) -> Printf.sprintf "%s %s" (Lexer.op2str op) (aux e depth)
    | Identifier id -> id
    | Number n -> string_of_int n in
  aux e 0

let rec parse_let_expr (): let_expr_tree parser = 
  let* _ = match_token (Lexer.Keyword Lexer.Let) in
  let* tt = next_token in
  match tt.token_type with
    Lexer.Identifier(x) ->
      let id = x in
      let* _ = match_token Lexer.Equal in
      let* e1 = expr() in
      let* _ = match_token (Lexer.Keyword Lexer.In) in
      let* e2 = expr() in
      return (Let (id, e1, e2))
  | _ -> stop
and parse_fun_expr (): fun_expr_tree parser =
  let* _ = match_token (Lexer.Keyword Lexer.Fun) in
  let* tt = next_token in
  match tt.token_type with
    Lexer.Identifier(x) ->
      let id = x in
      let* _ = match_token Lexer.Arrow in
      let* e = expr() in
      return (Fun (id, e))
  | _ -> stop
and parse_if_expr (): if_expr_tree parser =
  let* _ = match_token (Lexer.Keyword Lexer.If) in
  let* e1 = expr() in
  let* _ = match_token (Lexer.Keyword Lexer.Then) in
  let* e2 = expr() in
  let* _ = match_token (Lexer.Keyword Lexer.Else) in
  let* e3 = expr() in
  return (If (e1, e2, e3))
and parse_factor (): expr_tree parser =
  let* tt = peek_token in
  match tt.token_type with
  | Lexer.Identifier x ->
    let* _ = next_token in
    return (Identifier x)
  | Lexer.Digit x ->
    let* _ = next_token in
    return (Number (int_of_string x)) 
  | Lexer.LParen -> 
    let* _ = match_token Lexer.LParen in
    let* e = expr() in
    let* _ = match_token Lexer.RParen in
    return e
  | _ -> stop
and parse_call_expr (): expr_tree parser =
  let* e1 = parse_factor() in
  let rec aux e1 = 
    let* c = peek_token in
    match c.token_type with
    | Lexer.Identifier _ | Lexer.Digit _ | Lexer.LParen -> 
      let* e2 = parse_factor() in
      aux (CallExpr (Call (e1, e2)))
    | _ -> return e1 in
  aux e1
and parse_level1 (): expr_tree parser =
  let* e1 = parse_call_expr() in
  let rec aux e1 = 
    let* c = peek_token in
    match c.token_type with
    | Lexer.Op op when op = Lexer.Add || op = Lexer.Sub -> 
      let* _ = next_token in
      let* e2 = parse_call_expr() in
      aux (BinOpExpr (op, e1, e2))
    | _ -> return e1 in
  aux e1
and parse_level2 (): expr_tree parser =
  let* e1 = parse_level1() in
  let rec aux e1 = 
    let* c = peek_token in
    match c.token_type with
    | Lexer.Op op when op = Lexer.Mul || op = Lexer.Div || op = Lexer.Mod -> 
      let* _ = next_token in
      let* e2 = parse_level1() in
      aux (BinOpExpr (op, e1, e2))
    | _ -> return e1 in
  aux e1
and parse_level3 (): expr_tree parser =
  let* e1 = parse_level2() in
  let rec aux e1 = 
    let* c = peek_token in
    match c.token_type with
    | Lexer.Op op when op = Lexer.Pow -> 
      let* _ = next_token in
      let* e2 = parse_level3() in
      aux (BinOpExpr (op, e1, e2))
    | _ -> return e1 in
  aux e1
and expr (): expr_tree parser =
  let* e = (parse_let_expr() |> fmap (fun x -> LetExpr x)) <|> 
    (parse_fun_expr() |> fmap (fun x -> FunExpr x)) <|> 
    (parse_if_expr() |> fmap (fun x -> IfExpr x)) <|> parse_level3() in
  return e

let get_expr_tree_from_tokens (tokens: Lexer.token Seq.t): expr_tree option =
  let ntokens = Seq.filter (fun x -> 
    match x.token_type with 
    | Lexer.Comment(_) -> false 
    | _ -> true
  ) tokens in
  let ctx = { seq = ntokens; errors = [] } in
  match expr() ctx with
  | Some (e, _) -> Some e
  | None -> None

let%test "test get_expr_tree_from_tokens 1" =
  let tokens = Lexer.lex_tokens_seq "let x = 1 in\n  x" in
  let tokens = tokens |> Seq.map (fun (x,_) -> x) in
  match get_expr_tree_from_tokens tokens with
  | Some e -> expr2str e = "let x = 1 in\n  x"
  | None -> false