#include <catch2/catch_test_macros.hpp>
#include "dbc.hpp"
#include <iostream>
#include <vector>
#include "levelmanager.hpp"
#include "matrix.hpp"
#include "components.hpp"
#include <bitset>
#include <limits>

using namespace dbc;
using namespace components;

constexpr const int SCORE_MAX = std::numeric_limits<int>::max();

enum StateNames {
  ENEMY_IN_RANGE,
  ENEMY_DEAD,
  STATE_MAX
};

using GOAPState = std::bitset<STATE_MAX>;

bool is_subset(GOAPState& source, GOAPState& target) {

  GOAPState result = source & target;

  std::cout << "IS_SUBSET: source: " << source << " target: " << target << " result: " << result << " is it? " << (result == target) << std::endl;
  return result == target;
}

struct Action {
  std::string name;
  int cost = 0;

  std::unordered_map<StateNames, bool> preconds;
  std::unordered_map<StateNames, bool> effects;

  bool can_effect(GOAPState& state) {
    for(auto [name, setting] : preconds) {
      if(state[name] != setting) return false;
    }

    return true;
  }

  GOAPState apply_effect(GOAPState& state) {
    // RCR SUGGEST: state = (state & ~write_mask) | effect
    auto state_cp = state;

    for(auto [name, setting] : effects) {
      state_cp[name] = setting;
    }

    return state_cp;
  }

  bool operator==(const Action& other) const {
    return other.name == name;
  }
};

template<> struct std::hash<Action> {
  size_t operator()(const Action& p) const {
    return std::hash<std::string>{}(p.name);
  }
};

const Action FINAL_ACTION{"END", SCORE_MAX, {}, {}};

struct ActionState {
  Action action;
  GOAPState state;

  bool operator==(const ActionState& other) const {
    return other.action == action && other.state == state;
  }
};

template<> struct std::hash<ActionState> {
  size_t operator()(const ActionState& p) const {
      return std::hash<Action>{}(p.action) ^ std::hash<GOAPState>{}(p.state);
  }
};

using AStarPath = std::deque<Action>;

int distance_to_goal(GOAPState& from, GOAPState& to) {
  auto result = from ^ to;
  return result.count();
}


AStarPath reconstruct_path(std::unordered_map<Action, Action>& came_from, Action& current) {
  fmt::println(">> reconstruct path: {}", current.name);
  AStarPath total_path{current};
  int count = 0;

  while(came_from.contains(current) && count++ < 10) {
    current = came_from[current];
    if(current != FINAL_ACTION) {
      fmt::println("adding next action: {}", current.name);
      total_path.push_front(current);
    }
  }

  fmt::println("Exited reconstruct path.");
  return total_path;
}

inline int h(GOAPState& start, GOAPState& goal) {
  int result = distance_to_goal(start, goal);
  std::cout << "h on " << start << " and " << goal << " gives distance " << result << "\n";
  return result;
}

inline int d(GOAPState& start, GOAPState& goal) {
  int result = distance_to_goal(start, goal);
  std::cout << "d on " << start << " and " << goal << " gives distance " << result << "\n";
  return result;
}

inline ActionState find_lowest(std::unordered_map<ActionState, int>& open_set) {
  dbc::check(!open_set.empty(), "open set can't be empty in find_lowest");
  ActionState result;
  int lowest_score = SCORE_MAX;

  for(auto [as, score] : open_set) {
    fmt::println("### find_lowest: action={}, score={}", as.action.name, score);

    if(score < lowest_score) {
      lowest_score = score;
      result = as;
    }
  }

  fmt::println("<<< found lowest: action={}, score={}", result.action.name, lowest_score);
  return result;
}


// map is the list of possible actions
// start and goal are two world states
std::optional<AStarPath> plan_actions(std::vector<Action>& actions, GOAPState& start, GOAPState& goal) {
  std::unordered_map<ActionState, int> open_set;
  std::unordered_map<Action, Action> came_from;
  std::unordered_map<GOAPState, int> g_score;

  ActionState start_state{FINAL_ACTION, start};

  g_score[start] = 0;
  open_set[start_state] = g_score[start] + h(start, goal);

  while(!open_set.empty()) {
    fmt::println(">>>>>>>>>>>>>>>>>>>>>> TOP OF WHILE");
    auto current = find_lowest(open_set);

    if(is_subset(current.state, goal)) {
      return std::make_optional<AStarPath>(reconstruct_path(came_from, current.action));
    }

    open_set.erase(current);

    for(auto& neighbor_action : actions) {
      fmt::println("^^^ NEXT ACTION {}", neighbor_action.name);

      // calculate the GOAPState being current/neighbor
      if(!neighbor_action.can_effect(current.state)) {
        fmt::println("^^^ SKIP action {}", neighbor_action.name);
        continue;
      }

      auto neighbor = neighbor_action.apply_effect(current.state);
      int d_score = d(current.state, neighbor);
      int tentative_g_score = g_score[current.state] + d_score;
      int neighbor_g_score = g_score.contains(neighbor) ? g_score[neighbor] : SCORE_MAX;
      if(tentative_g_score < neighbor_g_score) {
        fmt::println("!!! NEW LOW SCORE::: SETTING {} with PARENT {}, tg_score={}, ng_score={}",
            neighbor_action.name, current.action.name,
            tentative_g_score, neighbor_g_score);

        came_from[neighbor_action] = current.action;

        g_score[neighbor] = tentative_g_score;
        // open_set gets the fScore
        ActionState neighbor_as{neighbor_action, neighbor};
        open_set[neighbor_as] = tentative_g_score + h(neighbor, goal);
      }
      fmt::println("^^^ END ACTION LOOP");
    }

    fmt::println("<<<<<<<<<<<<<<<<< END OF WHILE");
  }

  return std::nullopt;
}


TEST_CASE("worldstate works", "[goap]") {
  GOAPState goal;
  GOAPState start;
  std::vector<Action> actions;

  // start off enemy not dead and not in range
  start[ENEMY_DEAD] = false;
  start[ENEMY_IN_RANGE] = false;

  // end goal is enemy is dead
  goal[ENEMY_DEAD] = true;

  Action move_closer;
  move_closer.name = "move_closer";
  move_closer.cost = 10;
  move_closer.preconds[ENEMY_IN_RANGE] = false;
  move_closer.effects[ENEMY_IN_RANGE] = true;

  REQUIRE(move_closer.can_effect(start));
  auto after_move_state = move_closer.apply_effect(start);
  REQUIRE(start[ENEMY_IN_RANGE] == false);
  REQUIRE(after_move_state[ENEMY_IN_RANGE] == true);
  REQUIRE(after_move_state[ENEMY_DEAD] == false);
  // start is clean but after move is dirty
  REQUIRE(move_closer.can_effect(start));
  REQUIRE(!move_closer.can_effect(after_move_state));
  REQUIRE(distance_to_goal(start, after_move_state) == 1);

  Action kill_it;
  kill_it.name = "kill_it";
  kill_it.cost = 10;
  kill_it.preconds[ENEMY_IN_RANGE] = true;
  kill_it.preconds[ENEMY_DEAD] = false;
  kill_it.effects[ENEMY_DEAD] = true;

  REQUIRE(!kill_it.can_effect(start));
  REQUIRE(kill_it.can_effect(after_move_state));

  auto after_kill_state = kill_it.apply_effect(after_move_state);
  REQUIRE(!kill_it.can_effect(after_kill_state));
  REQUIRE(distance_to_goal(after_move_state, after_kill_state) == 1);

  actions.push_back(kill_it);
  actions.push_back(move_closer);

  REQUIRE(start != goal);
}

TEST_CASE("basic feature tests", "[goap]") {
  GOAPState goal;
  GOAPState start;
  std::vector<Action> actions;

  // start off enemy not dead and not in range
  start[ENEMY_DEAD] = false;
  start[ENEMY_IN_RANGE] = false;

  // end goal is enemy is dead
  goal[ENEMY_DEAD] = true;

  Action move_closer;
  move_closer.name = "move_closer";
  move_closer.cost = 10;
  move_closer.preconds[ENEMY_IN_RANGE] = false;
  move_closer.effects[ENEMY_IN_RANGE] = true;

  Action kill_it;
  kill_it.name = "kill_it";
  kill_it.cost = 10;
  kill_it.preconds[ENEMY_IN_RANGE] = true;
  kill_it.preconds[ENEMY_DEAD] = false;
  kill_it.effects[ENEMY_DEAD] = true;

  // order seems to matter which is wrong
  actions.push_back(kill_it);
  actions.push_back(move_closer);

  auto result = plan_actions(actions, start, goal);
  REQUIRE(result != std::nullopt);

  for(auto& action : *result) {
    fmt::println("ACTION: {}", action.name);
  }
}