#!/usr/bin/env elixir

defmodule EncounterDifficulty do
  @moduledoc """
  D&D 5e Encounter Difficulty Calculator

  Calculates encounter difficulty based on party composition and monster CRs.
  """

  # XP Thresholds by Character Level: %{level => {easy, medium, hard, deadly}}
  @xp_thresholds %{
    1 => {25, 50, 75, 100},
    2 => {50, 100, 150, 200},
    3 => {75, 150, 225, 400},
    4 => {125, 250, 375, 500},
    5 => {250, 500, 750, 1100},
    6 => {300, 600, 900, 1400},
    7 => {350, 750, 1100, 1700},
    8 => {450, 900, 1400, 2100},
    9 => {550, 1100, 1600, 2400},
    10 => {600, 1200, 1900, 2800},
    11 => {800, 1600, 2400, 3600},
    12 => {1000, 2000, 3000, 4500},
    13 => {1100, 2200, 3400, 5100},
    14 => {1250, 2500, 3800, 5700},
    15 => {1400, 2800, 4300, 6400},
    16 => {1600, 3200, 4800, 7200},
    17 => {2000, 3900, 5900, 8800},
    18 => {2100, 4200, 6300, 9500},
    19 => {2400, 4900, 7300, 10900},
    20 => {2800, 5700, 8500, 12700}
  }

  # XP by Challenge Rating
  @cr_xp %{
    "0" => 10,
    "1/8" => 25,
    "1/4" => 50,
    "1/2" => 100,
    "1" => 200,
    "2" => 450,
    "3" => 700,
    "4" => 1100,
    "5" => 1800,
    "6" => 2300,
    "7" => 2900,
    "8" => 3900,
    "9" => 5000,
    "10" => 5900,
    "11" => 7200,
    "12" => 8400,
    "13" => 10000,
    "14" => 11500,
    "15" => 13000,
    "16" => 15000,
    "17" => 18000,
    "18" => 20000,
    "19" => 22000,
    "20" => 25000,
    "21" => 33000,
    "22" => 41000,
    "23" => 50000,
    "24" => 62000,
    "25" => 75000,
    "26" => 90000,
    "27" => 105_000,
    "28" => 120_000,
    "29" => 135_000,
    "30" => 155_000
  }

  @doc "Get the encounter multiplier based on number of monsters."
  def get_multiplier(monster_count) do
    cond do
      monster_count == 1 -> 1.0
      monster_count == 2 -> 1.5
      monster_count in 3..6 -> 2.0
      monster_count in 7..10 -> 2.5
      monster_count in 11..14 -> 3.0
      true -> 4.0
    end
  end

  @doc "Normalize CR string for lookup."
  def parse_cr(cr_str) do
    cr_str = cr_str |> String.trim() |> String.downcase()

    cond do
      String.contains?(cr_str, "/") ->
        cr_str

      String.contains?(cr_str, ".") ->
        case Float.parse(cr_str) do
          {0.125, _} -> "1/8"
          {0.25, _} -> "1/4"
          {0.5, _} -> "1/2"
          _ -> cr_str
        end

      true ->
        cr_str
    end
  end

  @doc """
  Parse monster arguments in format 'count:CR' or 'CR' (assumes count=1).
  Examples: '2:1', '1/4', '3:1/2', '5'
  """
  def parse_monsters(monster_args) do
    Enum.map(monster_args, fn arg ->
      {count, cr_str} = parse_monster_arg(arg)
      cr = parse_cr(cr_str)

      unless Map.has_key?(@cr_xp, cr) do
        IO.puts(:stderr, "Error: Unknown CR '#{cr_str}'")
        System.halt(1)
      end

      {count, cr}
    end)
  end

  # TODO: Probably I can merge this with parse_count_value. I just need to consider the CR string case.
  defp parse_monster_arg(arg) do
    case String.split(arg, ":", parts: 2) do
      [count, cr] -> {String.to_integer(count), cr}
      [cr] -> {1, cr}
    end
  end

  @doc """
  Parse party arguments in format 'count:level' or 'level' (assumes count=1).
  Examples: '2:5', '3', '1:10'
  """
  def parse_party(party_args) do
    Enum.flat_map(party_args, fn arg ->
      {count, level} = parse_count_value(arg)

      if level < 1 or level > 20 do
        IO.puts(:stderr, "Error: Invalid level #{level}. Must be 1-20.")
        System.halt(1)
      end

      List.duplicate(level, count)
    end)
  end

  defp parse_count_value(arg) do
    case String.split(arg, ":", parts: 2) do
      [count, value] ->
        {String.to_integer(count), String.to_integer(value)}

      [value] ->
        {1, String.to_integer(value)}
    end
  end

  @doc "Calculate total party thresholds for each difficulty."
  def calculate_party_thresholds(party) do
    Enum.reduce(party, {0, 0, 0, 0}, fn level, {easy, medium, hard, deadly} ->
      {e, m, h, d} = @xp_thresholds[level]
      {easy + e, medium + m, hard + h, deadly + d}
    end)
  end

  @doc "Calculate encounter difficulty and XP."
  def calculate_encounter(monsters, party) do
    {total_monster_count, base_xp} =
      Enum.reduce(monsters, {0, 0}, fn {count, cr}, {tc, xp} ->
        {tc + count, xp + count * @cr_xp[cr]}
      end)

    multiplier = get_multiplier(total_monster_count)
    adjusted_xp = trunc(base_xp * multiplier)

    {easy, medium, hard, deadly} = calculate_party_thresholds(party)

    difficulty =
      cond do
        adjusted_xp >= deadly -> "Deadly"
        adjusted_xp >= hard -> "Hard"
        adjusted_xp >= medium -> "Medium"
        adjusted_xp >= easy -> "Easy"
        true -> "Trivial"
      end

    xp_per_character = div(base_xp, length(party))

    %{
      difficulty: difficulty,
      base_xp: base_xp,
      adjusted_xp: adjusted_xp,
      multiplier: multiplier,
      xp_per_character: xp_per_character,
      monster_count: total_monster_count,
      party_size: length(party),
      thresholds: %{
        easy: easy,
        medium: medium,
        hard: hard,
        deadly: deadly
      }
    }
  end

  def format_number(n) do
    n
    |> Integer.to_string()
    |> String.reverse()
    |> String.to_charlist()
    |> Enum.chunk_every(3)
    |> Enum.join(",")
    |> String.reverse()
  end

  def print_usage do
    IO.puts("""
    D&D 5e Encounter Difficulty Calculator

    Usage: elixir encounter_difficulty.exs -m MONSTERS -p PARTY

    Options:
      -m, --monsters  Monsters in format 'count:CR' or just 'CR'. Examples: 2:1, 1/4, 3:1/2
      -p, --party     Party members in format 'count:level' or just 'level'. Examples: 2:5, 3, 1:10
      -h, --help      Show this help message

    Examples:
      elixir encounter_difficulty.exs -m 2:1 1:2 -p 2:2 1:3
        2 CR 1 monsters and 1 CR 2 monster vs 2 level 2 and 1 level 3 characters

      elixir encounter_difficulty.exs -m 1/4 1/4 1/2 -p 4:1
        2 CR 1/4 and 1 CR 1/2 monster vs 4 level 1 characters

      elixir encounter_difficulty.exs --monsters 3:5 --party 4:10
        3 CR 5 monsters vs 4 level 10 characters
    """)
  end

  def parse_args(args) do
    case parse_args(args, %{monsters: [], party: [], current: nil}) do
      :help ->
        :help

      %{monsters: monsters, party: party} = acc ->
        %{acc | monsters: monsters, party: party}
    end
  end

  defp parse_args([], acc), do: acc

  defp parse_args(["-h" | _], _acc), do: :help
  defp parse_args(["--help" | _], _acc), do: :help

  defp parse_args(["-m" | rest], acc), do: parse_args(rest, %{acc | current: :monsters})
  defp parse_args(["--monsters" | rest], acc), do: parse_args(rest, %{acc | current: :monsters})

  defp parse_args(["-p" | rest], acc), do: parse_args(rest, %{acc | current: :party})
  defp parse_args(["--party" | rest], acc), do: parse_args(rest, %{acc | current: :party})

  defp parse_args([arg | rest], acc) do
    case acc.current do
      :monsters ->
        parse_args(rest, %{acc | monsters: [arg | acc.monsters]})

      :party ->
        parse_args(rest, %{acc | party: [arg | acc.party]})

      nil ->
        IO.puts(:stderr, "Error: Unexpected argument '#{arg}'")
        System.halt(1)
    end
  end

  def main(args) do
    case parse_args(args) do
      :help ->
        print_usage()

      %{monsters: [], party: _} ->
        IO.puts(:stderr, "Error: --monsters is required")
        print_usage()
        System.halt(1)

      %{monsters: _, party: []} ->
        IO.puts(:stderr, "Error: --party is required")
        print_usage()
        System.halt(1)

      %{monsters: monster_args, party: party_args} ->
        monsters = parse_monsters(monster_args)
        party = parse_party(party_args)
        result = calculate_encounter(monsters, party)

        IO.puts("")
        IO.puts(String.duplicate("=", 50))
        IO.puts("        D&D 5e ENCOUNTER DIFFICULTY")
        IO.puts(String.duplicate("=", 50))

        IO.puts("\nParty: #{result.party_size} characters")
        IO.puts("Monsters: #{result.monster_count} (x#{result.multiplier} multiplier)")

        IO.puts("\n--- Party XP Thresholds ---")
        IO.puts("  Easy:   #{format_number(result.thresholds.easy)} XP")
        IO.puts("  Medium: #{format_number(result.thresholds.medium)} XP")
        IO.puts("  Hard:   #{format_number(result.thresholds.hard)} XP")
        IO.puts("  Deadly: #{format_number(result.thresholds.deadly)} XP")

        IO.puts("\n--- Encounter ---")
        IO.puts("  Base XP:     #{format_number(result.base_xp)} XP")
        IO.puts("  Adjusted XP: #{format_number(result.adjusted_xp)} XP (for difficulty)")

        IO.puts("\n" <> String.duplicate("=", 50))
        IO.puts("  DIFFICULTY: #{String.upcase(result.difficulty)}")
        IO.puts(String.duplicate("=", 50))

        IO.puts("\n  Total XP Award:      #{format_number(result.base_xp)} XP")
        IO.puts("  XP per Character:    #{format_number(result.xp_per_character)} XP")
        IO.puts("")
    end
  end
end

EncounterDifficulty.main(System.argv())