open Lexer
open Expr

exception Parse_error of string * position

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

let token_name = function
  | Num _ -> "number"
  | Var _ -> "variable"
  | Ident _ -> "identifier"
  | Plus -> "'+'"
  | Minus -> "'-'"
  | Star -> "'*'"
  | Slash -> "'/'"
  | Caret -> "'^'"
  | LParen -> "'('"
  | RParen -> "')'"
  | Comma -> "','"
  | EOF -> "end of input"

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

let peek_pos state =
  if state.pos < List.length state.tokens then
    (List.nth state.tokens state.pos).start_pos
  else
    { line = 0; col = 0; offset = 0 }

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

let expect state tok =
  if peek state = tok then advance state
  else
    let pos = peek_pos state in
    raise (Parse_error (
      Printf.sprintf "expected %s but found %s at line %d, column %d"
        (token_name tok) (token_name (peek state)) pos.line pos.col,
      pos))

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;
      (match v with
       | "pi" | "π" -> SymConst Pi
       | "e" -> SymConst E
       | _ -> 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 "sinh" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Sinh arg
  | Ident "cosh" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Cosh arg
  | Ident "tanh" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Tanh arg
  | Ident "asin" | Ident "arcsin" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Asin arg
  | Ident "acos" | Ident "arccos" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Acos arg
  | Ident "atan" | Ident "arctan" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Atan arg
  | Ident "atan2" ->
      advance state;
      expect state LParen;
      let arg1 = parse_expr state in
      expect state Comma;
      let arg2 = parse_expr state in
      expect state RParen;
      Atan2 (arg1, arg2)
  | Ident "exp" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Exp arg
  | Ident "ln" | Ident "log" ->
      advance state;
      expect state LParen;
      let arg1 = parse_expr state in
      (match peek state with
       | Comma ->
           advance state;
           let arg2 = parse_expr state in
           expect state RParen;
           Log (arg1, arg2)
       | RParen ->
           advance state;
           Ln arg1
       | _ ->
           let pos = peek_pos state in
           raise (Parse_error (
             Printf.sprintf "expected ',' or ')' in log function at line %d, column %d"
               pos.line pos.col,
             pos)))
  | Ident "sqrt" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Sqrt arg
  | Ident "abs" ->
      advance state;
      expect state LParen;
      let arg = parse_expr state in
      expect state RParen;
      Abs arg
  | LParen ->
      advance state;
      let expr = parse_expr state in
      expect state RParen;
      expr
  | _ ->
      let pos = peek_pos state in
      raise (Parse_error (
        Printf.sprintf "unexpected %s at line %d, column %d"
          (token_name (peek state)) pos.line pos.col,
        pos))

let parse str =
  try
    let tokens = tokenize str in
    let state = { tokens; pos = 0 } in
    let expr = parse_expr state in
    if peek state = EOF then expr
    else
      let pos = peek_pos state in
      raise (Parse_error (
        Printf.sprintf "unexpected tokens after expression at line %d, column %d"
          pos.line pos.col,
        pos))
  with
  | Parse_error (msg, _) ->
      failwith msg
  | Failure msg ->
      failwith msg