#include "dbc.hpp"
#include "goap.hpp"

namespace ailol {
  bool is_subset(GOAPState& source, GOAPState& target) {
    GOAPState result = source & target;
    return result == target;
  }

  bool Action::can_effect(GOAPState& state) {
    return ((state & positive_preconds) == positive_preconds) &&
      ((state & negative_preconds) == ALL_ZERO);
  }

  GOAPState Action::apply_effect(GOAPState& state) {
    return (state | positive_effects) & ~negative_effects;
  }

  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) {
    AStarPath total_path{current};
    int count = 0;

    while(came_from.contains(current) && count++ < 10) {
      current = came_from.at(current);
      if(current != FINAL_ACTION) {
        total_path.push_front(current);
      }
    }

    return total_path;
  }

  inline int h(GOAPState& start, GOAPState& goal) {
    return distance_to_goal(start, goal);
  }

  inline int d(GOAPState& start, GOAPState& goal) {
    return distance_to_goal(start, goal);
  }

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

    for(auto& kv : open_set) {
      if(kv.second < lowest_score) {
        lowest_score = kv.second;
        result = &kv.first;
      }
    }

    return *result;
  }

  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()) {
      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) {
        // calculate the GOAPState being current/neighbor
        if(!neighbor_action.can_effect(current.state)) {
          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) {
          came_from.insert_or_assign(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);
        }
      }
    }

    return std::nullopt;
  }
}