open Lexer
open Expr

type parse_state = { tokens : token list; mutable pos : int }

let peek state =
  if state.pos < List.length state.tokens then
    List.nth state.tokens state.pos
  else EOF

let advance state =
  state.pos <- state.pos + 1

let expect state tok =
  if peek state = tok then advance state
  else failwith (Printf.sprintf "expected token")

let rec parse_expr state =
  parse_additive state

and parse_additive state =
  let left = ref (parse_multiplicative state) in
  while peek state = Plus || peek state = Minus do
    let op = peek state in
    advance state;
    let right = parse_multiplicative state in
    left := if op = Plus then Add (!left, right) else Sub (!left, right)
  done;
  !left

and parse_multiplicative state =
  let left = ref (parse_power state) in
  while peek state = Star || peek state = Slash do
    let op = peek state in
    advance state;
    let right = parse_power state in
    left := if op = Star then Mul (!left, right) else Div (!left, right)
  done;
  !left

and parse_power state =
  let left = parse_unary state in
  if peek state = Caret then begin
    advance state;
    let right = parse_power state in
    match left with
    | Var "e" -> Exp right
    | _ -> Pow (left, right)
  end else left

and parse_unary state =
  match peek state with
  | Minus ->
      advance state;
      Neg (parse_unary state)
  | _ -> parse_primary state

and parse_primary state =
  match peek state with
  | Num n ->
      advance state;
      Const n
  | Var v ->
      advance state;
      Var v
  | Ident "sin" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Sin arg
  | Ident "cos" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Cos arg
  | Ident "tan" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Tan arg
  | Ident "exp" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Exp arg
  | Ident "ln" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Ln arg
  | Ident "e" ->
      advance state;
      Const (exp 1.0)
  | LParen ->
      advance state;
      let expr = parse_expr state in
      expect state RParen;
      expr
  | _ -> failwith "unexpected token in primary"

let parse str =
  let tokens = tokenize str in
  let state = { tokens; pos = 0 } in
  let expr = parse_expr state in
  if peek state = EOF then expr
  else failwith "unexpected tokens after expression"