open Periodic
open Parser 

type element = 
  { symbol : string
  ; count : int
  }

type molecule =
  { elements : element list
  }

let unary_symbol_p : string parser =
  let* x = upper in
    String.make 1 x |> return

let binary_symbol_p : string parser =
  let* x = upper in
  let* x' = lower in
    String.make 1 x ^ String.make 1 x' |> return

let symbol_p : string parser =
  let* xs = binary_symbol_p <++ unary_symbol_p in
    match StringMap.find_opt xs periodic_table with
    | Some _ -> return xs
    | None   -> match element_find_closest xs with
                | [] -> fail "invalid atomic element"
                | xs -> List.fold_left (fun acc x -> acc ^ Printf.sprintf "%s " x) "" xs
                        |> Printf.sprintf "invalid atomic element. Did you mean any of the following: %s"
                        |> fail

let element_count_p : int parser =
  let* x = many1 digit in
    List.to_seq x
    |> String.of_seq
    |> int_of_string
    |> return

let element_p : element parser =
  let* s = symbol_p in
  let* n = element_count_p <++ return 1 in
    { symbol = s
    ; count = n
    } |> return

let molecule_p : molecule parser =
  let* xs = manyTill element_p eof in
    return { elements = xs }

let parse_molecule s : (molecule, error list) result =
  match make_input s |> molecule_p.run with
  | Ok []         -> failwith "TODO: make method total; empty list"
  | Ok ((x,_)::_) -> Ok x
  | Error xs      -> Error xs

let print_molecule (s: molecule) : unit =
  List.iter (fun x -> Printf.printf "%s, %i\n" x.symbol x.count) s.elements

let print_error (xs: error list) (s: string) : unit =
  List.iter (fun x -> Printf.printf "%s\n%*s\n%s\n\n" s x.column "^" x.content) xs

let molecule s : unit =
  match parse_molecule s with
  | Ok x     -> print_molecule x
  | Error xs -> print_error xs s

let () = molecule "C6Hi12O6"