#include "common/aoc.h"

namespace y2015::day22
{
  struct Player
  {
    int hp;
    int armor;
    int mana;
  };

  struct Monster
  {
    int hp;
    int dmg;
  };

  REGISTER_DAY(2015, Day22, Monster, int);

  REGISTER_TEST_EXAMPLE(2015, Day22, ExampleInput, 1, 212);
  REGISTER_TEST(2015, Day22, Input, 1, 900);
  REGISTER_TEST_EXAMPLE(2015, Day22, ExampleInput, 2, 212);
  REGISTER_TEST(2015, Day22, Input, 2, 1216);

  READ_INPUT(input)
  {
    Monster monster;
    input >> "Hit Points: " >> monster.hp >> "\n";
    input >> "Damage: " >> monster.dmg >> "\n";
    return monster;
  }

  struct Spell
  {
    int cost;
    int timer;
    int dmg;
    int armor;
    int manaGain;
    int heal;
  };

  struct State
  {
    Player player;
    Monster monster;
    std::vector<int> spellTimers;

    bool operator<(const State& other) const
    {
      if (player.hp != other.player.hp)
        return player.hp < other.player.hp;
      if (player.armor != other.player.armor)
        return player.armor < other.player.armor;
      if (player.mana != other.player.mana)
        return player.mana < other.player.mana;

      if (monster.hp != other.monster.hp)
        return monster.hp < other.monster.hp;
      if (monster.dmg != other.monster.dmg)
        return monster.dmg < other.monster.dmg;

      return spellTimers < other.spellTimers;
    }
  };

  // clang-format off
  //                                    cost  timer  dmg  armor  manaGain  heal
  const std::vector<Spell> spells{Spell{53,   1,     4,   0,     0,        0},
                                  Spell{73,   1,     2,   0,     0,        2},
                                  Spell{113,  6,     0,   7,     0,        0},
                                  Spell{173,  6,     3,   0,     0,        0},
                                  Spell{229,  5,     0,   0,     101,      0}};
  // clang-format-on

  void ApplyEffects(State& state)
  {
    for (int i = 0; i < state.spellTimers.size(); i++)
    {
      if (state.spellTimers[i] == 0)
        continue;

      const auto& spell = spells[i];
      state.spellTimers[i]--;
      state.monster.hp -= spell.dmg;
      state.player.hp += spell.heal;
      state.player.mana += spell.manaGain;
      if (state.spellTimers[i] == 0)
        state.player.armor -= spell.armor;
    }
  }

  std::pair<bool, int> PlayerTurn(State state, bool hardmode, std::map<State, std::pair<bool, int>>& memoization)
  {
    // Apply rules at start of turn
    if (hardmode)
      state.player.hp--;
    if (state.player.hp <= 0)
      return {false, 0};
    ApplyEffects(state);

    if (state.monster.hp <= 0)
      return {true, 0};

    auto it = memoization.find(state);
    if (it != memoization.end())
      return it->second;

    bool canWin = false;
    int minManaCost = std::numeric_limits<int>::max();
    for (int i = 0; i < spells.size(); i++)
    {
      // Spell is already active
      if (state.spellTimers[i] > 0)
        continue;

      const auto& spell = spells[i];
      if (state.player.mana < spell.cost)
        continue;

      // Cast spell
      State newState = state;
      newState.spellTimers[i] = spell.timer;
      newState.player.mana -= spell.cost;
      newState.player.armor += spell.armor;

      // Boss turn
      ApplyEffects(newState);
      if (newState.monster.hp <= 0)
      {
        if (spell.cost < minManaCost)
        {
          minManaCost = spell.cost;
          canWin = true;
        }
        continue;
      }
      newState.player.hp -= std::max(1, newState.monster.dmg - newState.player.armor);

      // Player turn
      auto result = PlayerTurn(newState, hardmode, memoization);
      if (result.first && result.second + spell.cost < minManaCost)
      {
        canWin = true;
        minManaCost = result.second + spell.cost;
      }
    }

    memoization[state] = {canWin, minManaCost};
    return {canWin, minManaCost};
  }

  OUTPUT1(input)
  {
    std::map<State, std::pair<bool, int>> memoization;
    return PlayerTurn(State{Player{50, 0, 500}, input, {0, 0, 0, 0, 0}}, false, memoization).second;
  }

  OUTPUT2(input)
  {
    std::map<State, std::pair<bool, int>> memoization;
    return PlayerTurn(State{Player{50, 0, 500}, input, {0, 0, 0, 0, 0}}, true, memoization).second;
  }
}