(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

let explode s = (* convert a string to a list of chars *)
  let rec exp i l =
   if i < 0 then l else exp (i - 1) (s.[i] :: l) in
  exp (String.length s - 1) []

let implode l = String.of_seq (List.to_seq l)

let (<<) f g x = f(g(x))
let (>>) f g x = g(f(x))

let rec take n l = (* take the first n elements *)
  match n with
  | 0 -> []
  | n ->
      match l with
      | [] -> []
      | x :: xs -> x :: (take (n-1) xs)

let rec drop n l = (* drop the first n elemnts *)
  match n with
  | 0 -> l
  | n -> drop (n-1) @@ List.tl l

let c2i = String.make 1 >> int_of_string

(* --- *)

type exp = Add of exp * exp | Mul of exp * exp | Num of int
type op = Plus | Mult

let rec string_of_exp = function
  | Num x -> string_of_int x
  | Mul (a, b) -> Printf.sprintf "(%s * %s)" (string_of_exp a) (string_of_exp b)
  | Add (a, b) -> Printf.sprintf "(%s + %s)" (string_of_exp a) (string_of_exp b)

let print_exp = print_string << string_of_exp
let print_exps = List.iter print_exp

let string_of_op = function | Plus -> "Plus" | Mult -> "Mult"
let print_op = print_string << string_of_op
let print_ops = List.iter print_op

let rec find_matching_paren count (x::xs) = (* find matching closing paren *)
  match x with
  | ')' -> if count = 0 then 0 else 1 + find_matching_paren (count-1) xs
  | '(' -> 1 + find_matching_paren (count+1) xs
  | n -> 1 + find_matching_paren count xs

let rec run_stack exps ops =
  if List.length exps = 1
  then List.hd exps
  else
    let (l :: r :: exps') = exps in
    let (op :: ops') = ops in
    match op with
    | Plus -> run_stack ((Add (l, r)) :: exps') ops'
    | Mult -> run_stack ((Mul (l, r)) :: exps') ops'

let read_exp line =
  let line' = line
    |> explode
    |> List.filter (fun x -> x != ' ')
  in

  let rec go exps ops = function
    | [] -> run_stack exps ops (* end of our string *)
    | '*' :: xs -> Mul (run_stack exps ops, go [] [] xs) (* "split" on Mul *)
    | '+' :: xs -> go exps (Plus :: ops) xs
    | '(' :: xs ->
        let matching = 1 + find_matching_paren 0 xs in
        let sub = take matching xs in
        let rest = drop matching xs in
        go ((go [] [] sub)::exps) ops rest
    | ')' :: xs -> go exps ops xs
    | num :: xs -> go ((Num (c2i num)) :: exps) ops xs
  in

  go [] [] line'

let rec eval = function
  | Num a -> a
  | Add (a, b) -> (eval a) + (eval b)
  | Mul (a, b) -> (eval a) * (eval b)

let () =
  let exps = read_file "day18.txt" |> List.map read_exp in

  let values = List.map eval exps in
  let sum = List.fold_left (+) 0 values in
  print_int sum;