#include "dbc.hpp"
#include "ai.hpp"

namespace ai {
  using namespace nlohmann;
  using namespace dbc;

  bool is_subset(State& source, State& target) {
    State result = source & target;
    return result == target;
  }

  void Action::needs(int name, bool val) {
    if(val) {
      $positive_preconds[name] = true;
      $negative_preconds[name] = false;
    } else {
      $negative_preconds[name] = true;
      $positive_preconds[name] = false;
    }
  }

  void Action::effect(int name, bool val) {
    if(val) {
      $positive_effects[name] = true;
      $negative_effects[name] = false;
    } else {
      $negative_effects[name] = true;
      $positive_effects[name] = false;
    }
  }


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

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

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

  Script reconstruct_path(std::unordered_map<Action, Action>& came_from, Action& current) {
    Script 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(State& start, State& goal) {
    return distance_to_goal(start, goal);
  }

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

  ActionState find_lowest(std::unordered_map<ActionState, int>& open_set) {
    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<Script> plan_actions(std::vector<Action>& actions, State& start, State& goal) {
    std::unordered_map<ActionState, int> open_set;
    std::unordered_map<Action, Action> came_from;
    std::unordered_map<State, 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<Script>(reconstruct_path(came_from, current.action));
      }

      open_set.erase(current);

      for(auto& neighbor_action : actions) {
        // calculate the State 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;
  }

  static AIManager AIMGR;
  static bool initialized = false;

  inline void validate_profile(nlohmann::json& profile) {
    for(auto& [name_key, value] : profile.items()) {
      check(value < STATE_MAX,
        fmt::format("profile field {} has value {} greater than STATE_MAX {}", (std::string)name_key, (int)value, STATE_MAX));
    }
  }

  Action config_action(nlohmann::json& profile, nlohmann::json& config) {
    check(config.contains("name"), "config_action: action config missing name");
    check(config.contains("cost"), "config_action: action config missing cost");

    validate_profile(profile);

    Action result(config["name"], config["cost"]);

    check(config.contains("needs"),
        fmt::format("config_action: no 'needs' field", result.$name));
    check(config.contains("effects"),
        fmt::format("config_action: no 'effects' field", result.$name));

    for(auto& [name_key, value] : config["needs"].items()) {
      check(profile.contains(name_key), fmt::format("config_action: profile does not have name {}", result.$name, name_key));
      int name = profile[name_key].template get<int>();

      result.needs(name, bool(value));
    }

    for(auto& [name_key, value] : config["effects"].items()) {
      check(profile.contains(name_key), fmt::format("config_action: profile does not have name {}", result.$name, name_key));

      int name = profile[name_key].template get<int>();

      result.effect(name, bool(value));
    }

    return result;
  }

  State config_state(nlohmann::json& profile, nlohmann::json& config) {
    State result;
    validate_profile(profile);

    for(auto& [name_key, value] : config.items()) {
      check(profile.contains(name_key), fmt::format("config_state: profile does not have name {}", name_key));

      int name = profile[name_key].template get<int>();
      result[name] = bool(value);
    }

    return result;
  }

  void init() {
    initialized = true;
    Config config("assets/ai.json");

    // profile specifies what keys (bitset indexes) are allowed
    // and how they map to the bitset of State
    AIMGR.profile = config["profile"];
    validate_profile(AIMGR.profile);

    // load all actions
    auto& actions = config["actions"];
    for(auto& action_vars : actions) {
      auto the_action = config_action(AIMGR.profile, action_vars);
      AIMGR.actions.insert_or_assign(the_action.$name, the_action);
    }

    // load all states
    auto& states = config["states"];
    for(auto& [name, state_vars] : states.items()) {
      auto the_state = config_state(AIMGR.profile, state_vars);
      AIMGR.states.insert_or_assign(name, the_state);
    }

    auto& scripts = config["scripts"];
    for(auto& [script_name, action_names] : scripts.items()) {
      std::vector<Action> the_script;
      for(auto name : action_names) {

        check(AIMGR.actions.contains(name),
            fmt::format("ai::init(): script {} uses action {} that doesn't exist",
              (std::string)script_name, (std::string)name));

        the_script.push_back(AIMGR.actions.at(name));
      }

      AIMGR.scripts.insert_or_assign(script_name, the_script);
    }
  }

  State load_state(std::string state_name) {
    check(initialized, "you forgot to initialize the AI first.");
    check(AIMGR.states.contains(state_name), fmt::format(
          "ai::load_state({}): state does not exist in config",
          state_name));

    return AIMGR.states.at(state_name);
  }

  Action load_action(std::string action_name) {
    check(initialized, "you forgot to initialize the AI first.");
    check(AIMGR.states.contains(action_name), fmt::format(
          "ai::load_action({}): action does not exist in config",
          action_name));
    return AIMGR.actions.at(action_name);
  }

  std::vector<Action> load_script(std::string script_name) {
    check(AIMGR.scripts.contains(script_name), fmt::format(
          "ai::load_script(): no script named {} configured", script_name));
    return AIMGR.scripts.at(script_name);
  }

  std::optional<Script> plan(std::string script_name, State start, State goal) {
    check(initialized, "you forgot to initialize the AI first.");
    auto script = load_script(script_name);
    return plan_actions(script, start, goal);
  }

  int state_id(std::string name) {
    check(AIMGR.profile.contains(name), fmt::format(
          "ai::state_id({}): id is not configured in profile",
          name));
    return AIMGR.profile.at(name);
  }
}