#include "tracker_state.h"

#include <fmt/core.h>
#include <hkutil/string.h>

#include <list>
#include <map>
#include <mutex>
#include <set>
#include <sstream>
#include <tuple>

#include "ap_state.h"
#include "game_data.h"
#include "logger.h"

namespace {

struct Requirements {
  bool disabled = false;

  std::set<int> doors;       // non-grouped, handles progressive
  std::set<int> panel_doors; // non-grouped, handles progressive
  std::set<int> items;       // all other items
  std::set<int> rooms;       // maybe
  bool mastery = false;      // maybe
  bool panel_hunt = false;   // maybe

  void Merge(const Requirements& rhs) {
    if (rhs.disabled) {
      return;
    }

    for (int id : rhs.doors) {
      doors.insert(id);
    }
    for (int id : rhs.panel_doors) {
      panel_doors.insert(id);
    }
    for (int id : rhs.items) {
      items.insert(id);
    }
    for (int id : rhs.rooms) {
      rooms.insert(id);
    }
    mastery = mastery || rhs.mastery;
    panel_hunt = panel_hunt || rhs.panel_hunt;
  }
};

class RequirementCalculator {
 public:
  void Reset() {
    doors_.clear();
    panels_.clear();
  }

  const Requirements& GetDoor(int door_id) {
    if (!doors_.count(door_id)) {
      Requirements requirements;
      const Door& door_obj = GD_GetDoor(door_id);

      if (door_obj.type == DoorType::kSunPainting) {
        if (!AP_IsPilgrimageEnabled()) {
          requirements.items.insert(door_obj.ap_item_id);
        } else {
          requirements.disabled = true;
        }
      } else if (door_obj.type == DoorType::kSunwarp) {
        switch (AP_GetSunwarpAccess()) {
          case kSUNWARP_ACCESS_NORMAL:
            // Do nothing.
            break;
          case kSUNWARP_ACCESS_DISABLED:
            requirements.disabled = true;
            break;
          case kSUNWARP_ACCESS_UNLOCK:
            requirements.items.insert(door_obj.group_ap_item_id);
            break;
          case kSUNWARP_ACCESS_INDIVIDUAL:
          case kSUNWARP_ACCESS_PROGRESSIVE:
            requirements.doors.insert(door_obj.id);
            break;
        }
      } else if (AP_GetDoorShuffleMode() != kDOORS_MODE || door_obj.skip_item) {
        requirements.rooms.insert(door_obj.room);

        for (int panel_id : door_obj.panels) {
          const Requirements& panel_reqs = GetPanel(panel_id);
          requirements.Merge(panel_reqs);
        }
      } else if (AP_AreDoorsGrouped() && !door_obj.group_name.empty()) {
        requirements.items.insert(door_obj.group_ap_item_id);
      } else {
        requirements.doors.insert(door_obj.id);
      }

      doors_[door_id] = requirements;
    }

    return doors_[door_id];
  }

  const Requirements& GetPanel(int panel_id) {
    if (!panels_.count(panel_id)) {
      Requirements requirements;
      const Panel& panel_obj = GD_GetPanel(panel_id);

      requirements.rooms.insert(panel_obj.room);

      if (panel_obj.name == "THE MASTER") {
        requirements.mastery = true;
      }

      if ((panel_obj.name == "ANOTHER TRY" || panel_obj.name == "LEVEL 2") &&
          AP_GetLevel2Requirement() > 1) {
        requirements.panel_hunt = true;
      }

      for (int room_id : panel_obj.required_rooms) {
        requirements.rooms.insert(room_id);
      }

      for (int door_id : panel_obj.required_doors) {
        const Requirements& door_reqs = GetDoor(door_id);
        requirements.Merge(door_reqs);
      }

      for (int panel_id : panel_obj.required_panels) {
        const Requirements& panel_reqs = GetPanel(panel_id);
        requirements.Merge(panel_reqs);
      }

      if (AP_IsColorShuffle()) {
        for (LingoColor color : panel_obj.colors) {
          requirements.items.insert(GD_GetItemIdForColor(color));
        }
      }
      
      if (panel_obj.panel_door != -1 &&
          AP_GetDoorShuffleMode() == kPANELS_MODE) {
        const PanelDoor& panel_door_obj = GD_GetPanelDoor(panel_obj.panel_door);

        if (panel_door_obj.group_ap_item_id != -1 && AP_AreDoorsGrouped()) {
          requirements.items.insert(panel_door_obj.group_ap_item_id);
        } else {
          requirements.panel_doors.insert(panel_obj.panel_door);
        }
      }

      panels_[panel_id] = requirements;
    }

    return panels_[panel_id];
  }

 private:
  std::map<int, Requirements> doors_;
  std::map<int, Requirements> panels_;
};

struct TrackerState {
  std::map<int, bool> reachability;
  std::set<int> reachable_doors;
  std::set<int> reachable_paintings;
  std::mutex reachability_mutex;
  RequirementCalculator requirements;
  std::map<int, std::map<std::string, bool>> door_reports;
  bool pilgrimage_doable = false;
};

enum Decision { kYes, kNo, kMaybe };

TrackerState& GetState() {
  static TrackerState* instance = new TrackerState();
  return *instance;
}

class StateCalculator;

struct StateCalculatorOptions {
  int start;
  bool pilgrimage = false;
  StateCalculator* parent = nullptr;
};

class StateCalculator {
 public:
  StateCalculator() = default;

  explicit StateCalculator(StateCalculatorOptions options)
      : options_(options) {}

  void Calculate() {
    painting_mapping_ = AP_GetPaintingMapping();
    checked_paintings_ = AP_GetCheckedPaintings();
    sunwarp_mapping_ = AP_GetSunwarpMapping();

    std::list<int> panel_boundary;
    std::list<int> painting_boundary;
    std::list<Exit> flood_boundary;
    flood_boundary.push_back(
        {.source_room = -1, .destination_room = options_.start});

    bool reachable_changed = true;
    while (reachable_changed) {
      reachable_changed = false;

      std::list<Exit> new_boundary;

      std::list<int> new_panel_boundary;
      for (int panel_id : panel_boundary) {
        if (solveable_panels_.count(panel_id)) {
          continue;
        }

        Decision panel_reachable = IsPanelReachable(panel_id);
        if (panel_reachable == kYes) {
          solveable_panels_.insert(panel_id);
          reachable_changed = true;
        } else if (panel_reachable == kMaybe) {
          new_panel_boundary.push_back(panel_id);
        }
      }

      std::list<int> new_painting_boundary;
      for (int painting_id : painting_boundary) {
        if (reachable_paintings_.count(painting_id)) {
          continue;
        }

        Decision painting_reachable = IsPaintingReachable(painting_id);
        if (painting_reachable == kYes) {
          reachable_paintings_.insert(painting_id);
          reachable_changed = true;

          PaintingExit cur_painting = GD_GetPaintingExit(painting_id);
          if (painting_mapping_.count(cur_painting.internal_id) &&
              checked_paintings_.count(cur_painting.internal_id)) {
            Exit painting_exit;
            PaintingExit target_painting =
                GD_GetPaintingExit(GD_GetPaintingByName(
                    painting_mapping_.at(cur_painting.internal_id)));
            painting_exit.source_room = cur_painting.room;
            painting_exit.destination_room = target_painting.room;
            painting_exit.type = EntranceType::kPainting;

            new_boundary.push_back(painting_exit);
          }
        } else if (painting_reachable == kMaybe) {
          new_painting_boundary.push_back(painting_id);
        }
      }

      for (const Exit& room_exit : flood_boundary) {
        if (reachable_rooms_.count(room_exit.destination_room)) {
          continue;
        }

        bool valid_transition = false;

        Decision exit_usable = IsExitUsable(room_exit);
        if (exit_usable == kYes) {
          valid_transition = true;
        } else if (exit_usable == kMaybe) {
          new_boundary.push_back(room_exit);
        }

        if (valid_transition) {
          reachable_rooms_.insert(room_exit.destination_room);
          reachable_changed = true;

#ifndef NDEBUG
          std::list<int> room_path = paths_[room_exit.source_room];
          room_path.push_back(room_exit.destination_room);
          paths_[room_exit.destination_room] = room_path;
#endif

          const Room& room_obj = GD_GetRoom(room_exit.destination_room);
          for (const Exit& out_edge : room_obj.exits) {
            if (out_edge.type == EntranceType::kPainting &&
                AP_IsPaintingShuffle()) {
              continue;
            }

            if (out_edge.type == EntranceType::kSunwarp &&
                AP_IsSunwarpShuffle()) {
              continue;
            }

            new_boundary.push_back(out_edge);
          }

          if (AP_IsPaintingShuffle()) {
            for (int out_edge : room_obj.paintings) {
              new_painting_boundary.push_back(out_edge);
            }
          }

          if (AP_IsSunwarpShuffle()) {
            for (int index : room_obj.sunwarps) {
              if (sunwarp_mapping_.count(index)) {
                const SunwarpMapping& sm = sunwarp_mapping_.at(index);

                new_boundary.push_back(
                    {.source_room = room_exit.destination_room,
                     .destination_room = GD_GetRoomForSunwarp(sm.exit_index),
                     .door = GD_GetSunwarpDoors().at(sm.dots - 1),
                     .type = EntranceType::kSunwarp});
              }
            }
          }

          if (AP_HasEarlyColorHallways() && room_obj.name == "Starting Room") {
            new_boundary.push_back(
                {.source_room = room_exit.destination_room,
                 .destination_room = GD_GetRoomByName("Color Hallways"),
                 .type = EntranceType::kPainting});
          }

          if (AP_IsPilgrimageEnabled()) {
            int pilgrimage_start_id = GD_GetRoomByName("Hub Room");
            if (AP_IsSunwarpShuffle()) {
              for (const auto& [start_index, mapping] : sunwarp_mapping_) {
                if (mapping.dots == 1) {
                  pilgrimage_start_id = GD_GetRoomForSunwarp(start_index);
                }
              }
            }

            if (room_exit.destination_room == pilgrimage_start_id) {
              new_boundary.push_back(
                  {.source_room = room_exit.destination_room,
                   .destination_room = GD_GetRoomByName("Pilgrim Antechamber"),
                   .type = EntranceType::kPilgrimage});
            }
          } else {
            if (room_obj.name == "Starting Room") {
              new_boundary.push_back(
                  {.source_room = room_exit.destination_room,
                   .destination_room = GD_GetRoomByName("Pilgrim Antechamber"),
                   .door =
                       GD_GetDoorByName("Pilgrim Antechamber - Sun Painting"),
                   .type = EntranceType::kPainting});
            }
          }

          for (int panel_id : room_obj.panels) {
            new_panel_boundary.push_back(panel_id);
          }
        }
      }

      flood_boundary = new_boundary;
      panel_boundary = new_panel_boundary;
      painting_boundary = new_painting_boundary;
    }

    // Now that we know the full reachable area, let's make sure all doors are
    // evaluated.
    for (const Door& door : GD_GetDoors()) {
      int discard = IsDoorReachable(door.id);
    }
  }

  const std::set<int>& GetReachableRooms() const { return reachable_rooms_; }

  const std::map<int, Decision>& GetDoorDecisions() const {
    return door_decisions_;
  }

  const std::set<int>& GetSolveablePanels() const { return solveable_panels_; }

  const std::set<int>& GetReachablePaintings() const {
    return reachable_paintings_;
  }

  const std::map<int, std::map<std::string, bool>>& GetDoorReports() const {
    return door_report_;
  }

  bool IsPilgrimageDoable() const { return pilgrimage_doable_; }

  std::string GetPathToRoom(int room_id) const {
    if (!paths_.count(room_id)) {
      return "";
    }

    const std::list<int>& path = paths_.at(room_id);
    std::vector<std::string> room_names;
    for (int room_id : path) {
      room_names.push_back(GD_GetRoom(room_id).name);
    }
    return hatkirby::implode(room_names, " -> ");
  }

 private:
  template <typename T>
  Decision IsNonGroupedDoorReachable(const T& door_obj) {
    bool has_item = AP_HasItem(door_obj.ap_item_id);

    if (!has_item) {
      for (const ProgressiveRequirement& prog_req : door_obj.progressives) {
        if (AP_HasItem(prog_req.ap_item_id, prog_req.quantity)) {
          has_item = true;
          break;
        }
      }
    }

    return has_item ? kYes : kNo;
  }

  Decision AreRequirementsSatisfied(
      const Requirements& reqs, std::map<std::string, bool>* report = nullptr) {
    if (reqs.disabled) {
      return kNo;
    }

    Decision final_decision = kYes;

    for (int door_id : reqs.doors) {
      const Door& door_obj = GD_GetDoor(door_id);
      Decision decision = IsNonGroupedDoorReachable(door_obj);

      if (report) {
        (*report)[door_obj.item_name] = (decision == kYes);
      }

      if (decision != kYes) {
        final_decision = decision;
      }
    }

    for (int panel_door_id : reqs.panel_doors) {
      const PanelDoor& panel_door_obj = GD_GetPanelDoor(panel_door_id);
      Decision decision = IsNonGroupedDoorReachable(panel_door_obj);

      if (report) {
        (*report)[AP_GetItemName(panel_door_obj.ap_item_id)] =
            (decision == kYes);
      }

      if (decision != kYes) {
        final_decision = decision;
      }
    }

    for (int item_id : reqs.items) {
      bool has_item = AP_HasItem(item_id);
      if (report) {
        (*report)[AP_GetItemName(item_id)] = has_item;
      }

      if (!has_item) {
        final_decision = kNo;
      }
    }

    for (int room_id : reqs.rooms) {
      bool reachable = reachable_rooms_.count(room_id);

      if (report) {
        std::string report_name = "Reach \"" + GD_GetRoom(room_id).name + "\"";
        (*report)[report_name] = reachable;
      }

      if (!reachable && final_decision != kNo) {
        final_decision = kMaybe;
      }
    }

    if (reqs.mastery) {
      int achievements_accessible = 0;

      for (int achieve_id : GD_GetAchievementPanels()) {
        if (solveable_panels_.count(achieve_id)) {
          achievements_accessible++;

          if (achievements_accessible >= AP_GetMasteryRequirement()) {
            break;
          }
        }
      }

      bool can_mastery =
          (achievements_accessible >= AP_GetMasteryRequirement());
      if (report) {
        (*report)["Mastery"] = can_mastery;
      }

      if (!can_mastery && final_decision != kNo) {
        final_decision = kMaybe;
      }
    }

    if (reqs.panel_hunt) {
      int counting_panels_accessible = 0;

      for (int solved_panel_id : solveable_panels_) {
        const Panel& solved_panel = GD_GetPanel(solved_panel_id);

        if (!solved_panel.non_counting) {
          counting_panels_accessible++;
        }
      }

      bool can_level2 =
          (counting_panels_accessible >= AP_GetLevel2Requirement() - 1);
      if (report) {
        std::string report_name =
            std::to_string(AP_GetLevel2Requirement()) + " Panels";
        (*report)[report_name] = can_level2;
      }

      if (!can_level2 && final_decision != kNo) {
        final_decision = kMaybe;
      }
    }

    return final_decision;
  }

  Decision IsDoorReachable_Helper(int door_id) {
    if (door_report_.count(door_id)) {
      door_report_[door_id].clear();
    } else {
      door_report_[door_id] = {};
    }

    return AreRequirementsSatisfied(GetState().requirements.GetDoor(door_id),
                                    &door_report_[door_id]);
  }

  Decision IsDoorReachable(int door_id) {
    if (options_.parent) {
      return options_.parent->IsDoorReachable(door_id);
    }

    if (door_decisions_.count(door_id)) {
      return door_decisions_.at(door_id);
    }

    Decision result = IsDoorReachable_Helper(door_id);
    if (result != kMaybe) {
      door_decisions_[door_id] = result;
    }

    return result;
  }

  Decision IsPanelReachable(int panel_id) {
    return AreRequirementsSatisfied(GetState().requirements.GetPanel(panel_id));
  }

  Decision IsPaintingReachable(int painting_id) {
    const PaintingExit& painting = GD_GetPaintingExit(painting_id);
    if (painting.door) {
      return IsDoorReachable(*painting.door);
    }

    return kYes;
  }

  Decision IsExitUsable(const Exit& room_exit) {
    if (room_exit.type == EntranceType::kPilgrimage) {
      if (options_.pilgrimage || !AP_IsPilgrimageEnabled()) {
        return kNo;
      }

      if (AP_GetSunwarpAccess() != kSUNWARP_ACCESS_NORMAL) {
        for (int door_id : GD_GetSunwarpDoors()) {
          Decision sub_decision = IsDoorReachable(door_id);
          if (sub_decision != kYes) {
            return sub_decision;
          }
        }
      }

      std::vector<std::tuple<int, int>> pilgrimage_pairs;
      if (AP_IsSunwarpShuffle()) {
        pilgrimage_pairs = std::vector<std::tuple<int, int>>(5);

        for (const auto& [start_index, mapping] : sunwarp_mapping_) {
          if (mapping.dots > 1) {
            std::get<1>(pilgrimage_pairs[mapping.dots - 2]) = start_index;
          }
          if (mapping.dots < 6) {
            std::get<0>(pilgrimage_pairs[mapping.dots - 1]) =
                mapping.exit_index;
          }
        }
      } else {
        pilgrimage_pairs = {{6, 1}, {8, 3}, {9, 4}, {10, 5}};
      }

      for (const auto& [from_sunwarp, to_sunwarp] : pilgrimage_pairs) {
        StateCalculator pilgrimage_calculator(
            {.start = GD_GetRoomForSunwarp(from_sunwarp),
             .pilgrimage = true,
             .parent = this});
        pilgrimage_calculator.Calculate();

        if (!pilgrimage_calculator.GetReachableRooms().count(
                GD_GetRoomForSunwarp(to_sunwarp))) {
          return kMaybe;
        }
      }

      pilgrimage_doable_ = true;

      return kYes;
    }

    if (options_.pilgrimage) {
      if (room_exit.type == EntranceType::kWarp ||
          room_exit.type == EntranceType::kSunwarp) {
        return kNo;
      }
      if (room_exit.type == EntranceType::kCrossroadsRoofAccess &&
          !AP_DoesPilgrimageAllowRoofAccess()) {
        return kNo;
      }
      if (room_exit.type == EntranceType::kPainting &&
          !AP_DoesPilgrimageAllowPaintings()) {
        return kNo;
      }
    }

    if (room_exit.type == EntranceType::kSunwarp) {
      if (AP_GetSunwarpAccess() == kSUNWARP_ACCESS_NORMAL) {
        return kYes;
      } else if (AP_GetSunwarpAccess() == kSUNWARP_ACCESS_DISABLED) {
        return kNo;
      }
    }

    if (room_exit.door.has_value()) {
      return IsDoorReachable(*room_exit.door);
    }

    return kYes;
  }

  StateCalculatorOptions options_;

  std::set<int> reachable_rooms_;
  std::map<int, Decision> door_decisions_;
  std::set<int> solveable_panels_;
  std::set<int> reachable_paintings_;
  std::map<int, std::map<std::string, bool>> door_report_;
  bool pilgrimage_doable_ = false;

  std::map<int, std::list<int>> paths_;

  std::map<std::string, std::string> painting_mapping_;
  std::set<std::string> checked_paintings_;
  std::map<int, SunwarpMapping> sunwarp_mapping_;
};

}  // namespace

void ResetReachabilityRequirements() {
  std::lock_guard reachability_guard(GetState().reachability_mutex);
  GetState().requirements.Reset();
}

void RecalculateReachability() {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  StateCalculator state_calculator({.start = GD_GetRoomByName("Menu")});
  state_calculator.Calculate();

  const std::set<int>& reachable_rooms = state_calculator.GetReachableRooms();
  const std::set<int>& solveable_panels = state_calculator.GetSolveablePanels();

  std::map<int, bool> new_reachability;
  for (const MapArea& map_area : GD_GetMapAreas()) {
    for (size_t section_id = 0; section_id < map_area.locations.size();
         section_id++) {
      const Location& location_section = map_area.locations.at(section_id);
      bool reachable = reachable_rooms.count(location_section.room);
      if (reachable) {
        for (int panel_id : location_section.panels) {
          reachable &= (solveable_panels.count(panel_id) == 1);
        }
      }

      new_reachability[location_section.ap_location_id] = reachable;
    }
  }

  std::set<int> new_reachable_doors;
  for (const auto& [door_id, decision] : state_calculator.GetDoorDecisions()) {
    if (decision == kYes) {
      new_reachable_doors.insert(door_id);
    }
  }

  std::set<int> reachable_paintings = state_calculator.GetReachablePaintings();
  std::map<int, std::map<std::string, bool>> door_reports =
      state_calculator.GetDoorReports();
  
  std::swap(GetState().reachability, new_reachability);
  std::swap(GetState().reachable_doors, new_reachable_doors);
  std::swap(GetState().reachable_paintings, reachable_paintings);
  std::swap(GetState().door_reports, door_reports);
  GetState().pilgrimage_doable = state_calculator.IsPilgrimageDoable();
}

bool IsLocationReachable(int location_id) {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  if (GetState().reachability.count(location_id)) {
    return GetState().reachability.at(location_id);
  } else {
    return false;
  }
}

bool IsDoorOpen(int door_id) {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  return GetState().reachable_doors.count(door_id);
}

bool IsPaintingReachable(int painting_id) {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  return GetState().reachable_paintings.count(painting_id);
}

const std::map<std::string, bool>& GetDoorRequirements(int door_id) {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  return GetState().door_reports[door_id];
}

bool IsPilgrimageDoable() {
  std::lock_guard reachability_guard(GetState().reachability_mutex);

  return GetState().pilgrimage_doable;
}