let read_lines path =
  let lines = ref [] in
  let chan = open_in path in
  try
    while true; do
      lines := input_line chan :: !lines
    done; !lines
  with End_of_file ->
    close_in chan;
    List.rev !lines

type packet_body =
  | Simple of int
  | Compound of packet list

and packet = {
  version: int;
  packet_type: packet_type;
  body: packet_body;
}

and packet_type =
  | Sum
  | Product
  | Minimum
  | Maximum
  | Literal
  | Greater
  | Lesser
  | Equal

let hexadecimal_digit_to_binary = function
  | '0' -> "0000"
  | '1' -> "0001"
  | '2' -> "0010"
  | '3' -> "0011"
  | '4' -> "0100"
  | '5' -> "0101"
  | '6' -> "0110"
  | '7' -> "0111"
  | '8' -> "1000"
  | '9' -> "1001"
  | 'A' -> "1010"
  | 'B' -> "1011"
  | 'C' -> "1100"
  | 'D' -> "1101"
  | 'E' -> "1110"
  | 'F' -> "1111"
  | _ -> assert false

let parse_packet_type = function
  | 0 -> Sum
  | 1 -> Product
  | 2 -> Minimum
  | 3 -> Maximum
  | 4 -> Literal
  | 5 -> Greater
  | 6 -> Lesser
  | 7 -> Equal
  | _ -> assert false

let hexadecimal_to_binary hex =
  let buffer = Buffer.create ((String.length hex) * 4) in
  let expand_digit digit = Buffer.add_string buffer (hexadecimal_digit_to_binary digit) in
  String.iter expand_digit hex;
  Buffer.contents buffer

let parse input =
  let data = hexadecimal_to_binary input in

  let rec binary_to_decimal ?(init = 0) start length =
    let stop = start + length in
    let rec parse acc position =
      if position < stop
      then parse (acc * 2 + (Char.code data.[position] - (Char.code '0'))) (position + 1)
      else acc
    in parse init start

  and parse_bytes start length = (binary_to_decimal start length, start + length)

  and parse_literal start =
    let rec aux position value =
      let value = binary_to_decimal ~init:value (position + 1) 4 in
      let next_pos = position + 5 in
      if data.[position] == '1'
      then aux next_pos value
      else (value, next_pos)
    in aux start 0

  and parse_n_subpackets start =
    let rec aux acc position n =
      match n with
      | 0 -> (List.rev acc, position)
      | _ -> let (packet, next) = parse_packet position
        in aux (packet :: acc) next (n - 1)
    in aux [] (start + 11) (binary_to_decimal start 11)

  and parse_packet position =
    let (version, next) = parse_bytes position 3 in
    let (packet_t, next) = parse_bytes next 3 in
    let (body, next) = match packet_t with
      | 4 -> let (value, next) = parse_literal next in (Simple value, next)
      | _ -> let (value, next) = parse_subpacket next in (Compound value, next)
    in ({ version = version; body = body; packet_type = parse_packet_type packet_t }, next)

  and parse_subpacket start =
    let (length_t, next) = parse_bytes start 1 in
    match length_t with
    | 0 -> parse_n_bytes_of_subpackets next
    | 1 -> parse_n_subpackets next
    | _ -> assert false

  and parse_n_bytes_of_subpackets start =
    let (count, next) = parse_bytes start 15 in
    let limit = count + next in
    let rec aux acc position =
      if position >= limit
      then (List.rev acc, position)
      else let (packet, next) = parse_packet position
        in aux (packet :: acc) next
    in aux [] next
  in

  let (packet, _) = parse_packet 0 in packet

let rec packet_version_sum packet =
  match packet.body with
  | Simple _ -> packet.version
  | Compound subpackets -> packet.version + (subpackets |> List.map packet_version_sum |> List.fold_left (+) 0)

let rec evaluate_packet packet =
  match packet.packet_type with
  | Literal -> (match packet.body with Simple x -> x | _ -> assert false)
  | Sum -> fold_left_subpackets (+) packet
  | Product -> fold_left_subpackets ( * ) packet
  | Minimum -> fold_left_subpackets min packet
  | Maximum -> fold_left_subpackets max packet
  | Greater -> subpacket_compare (>) packet
  | Lesser -> subpacket_compare (<) packet
  | Equal -> subpacket_compare (=) packet

and evaluate_subpackets packet = packet |> subpackets |> List.map evaluate_packet

and fold_left_subpackets f packet =
  let subs = evaluate_subpackets packet
  in List.fold_left f (List.hd subs) (List.tl subs)

and subpacket_compare f packet =
  match evaluate_subpackets packet with
  | a :: b :: [] -> if f a b then 1 else 0
  | _ -> assert false

and subpackets packet =
  match packet.body with
  | Compound x -> x
  | _ -> assert false

let part1 = packet_version_sum
let part2 = evaluate_packet

let packet = read_lines Sys.argv.(1) |> List.hd |> parse
let () = Printf.printf "Part 1: %d\n" (part1 packet)
let () = Printf.printf "Part 2: %d\n" (part2 packet)