defmodule Day05 do
  def parse_integer_list(list, split) do
    list
    |> String.split(split, trim: true)
    |> Enum.map(&String.to_integer/1)
  end

  def parse_data(data) do
    [rules, jobs] = data
    |> String.split("\n\n", trim: true)

    rules = rules
    |> String.split("\n", trim: true)
    |> Enum.map(&parse_integer_list(&1, "|"))
    |> Enum.group_by(&hd/1, &hd(tl(&1)))
    |> Enum.into(%{}, fn {k, v} -> {k, MapSet.new(v)} end)

    jobs = jobs
    |> String.split("\n", trim: true)
    |> Enum.map(&parse_integer_list(&1, ","))

    {rules, jobs}
  end

  def part1({rules, jobs}) do
    jobs
    |> Enum.filter(&job_okay?(&1, rules))
    |> Enum.map(&middle_page/1)
    |> Enum.sum()
  end

  def job_okay?(job, rules) do
    case job do
      [] ->
        true
      [page | rest] ->
        page_okay?(page, rest, rules) and
          job_okay?(rest, rules)
    end
  end

  def page_okay?(page, rest, rules) do
    rest
    |> Enum.all?(&Map.get(rules, &1, %MapSet{}) |> MapSet.member?(page) == false)
  end

  def middle_page(job) do
    job |> Enum.at(job |> length() |> div(2))
  end

  def part2({rules, jobs}) do
    jobs
    |> Enum.reject(&job_okay?(&1, rules))
    |> Enum.map(&fix_job(&1, rules))
    |> Enum.map(&middle_page/1)
    |> Enum.sum()
  end

  def fix_job(job, rules, acc \\ []) do
    case job do
      [] ->
        Enum.reverse(acc)
      job ->
        {l, r} = Enum.split_with(job, &page_okay?(&1, job, rules) == false)
        [next, rest] = [hd(r), l ++ tl(r)]
        fix_job(rest, rules, [next | acc])
    end
  end
end

data = IO.read(:stdio, :eof) |> Day05.parse_data()

{time1 , ans1} = :timer.tc(fn -> Day05.part1(data) end)
IO.puts("Time  : #{time1 / 1000000}")
IO.puts("Answer: #{ans1}")

{time2 , ans2} = :timer.tc(fn -> Day05.part2(data) end)
IO.puts("Time  : #{time2 / 1000000}")
IO.puts("Answer: #{ans2}")