#include "state.hpp"
#include "presets.hpp"
#include "tracy.hpp"

#include <mutex>
#include <raymath.h>

auto StateManager::LoadPreset(int preset) -> void {
  current_state = generators[preset]();
  previous_state = current_state;
  ClearGraph();
  current_preset = preset;
  edited = false;
}

auto StateManager::ResetState() -> void {
  current_state = generators[current_preset]();
  previous_state = current_state;
  if (edited) {
    // We also need to clear the graph in case the state has been edited
    // because the graph could contain states that are impossible to reach now.
    ClearGraph();
    edited = false;
  }
}

auto StateManager::PreviousPreset() -> void {
  LoadPreset((generators.size() + current_preset - 1) % generators.size());
}

auto StateManager::NextPreset() -> void {
  LoadPreset((current_preset + 1) % generators.size());
}

auto StateManager::FillGraph() -> void {
  ClearGraph();

  std::pair<std::unordered_set<State>, std::vector<std::pair<State, State>>>
      closure = current_state.Closure();

  physics.AddMassSpringsCmd(closure.first, closure.second);

  // TODO: We have only dispatched the commands, the states won't be downloaded
  //       when calling this... Make FindWinningStates() another command?
  //       Or recalculate whenever masses.size() changes?
  // FindWinningStates();
}

auto StateManager::UpdateGraph() -> void {
  if (previous_state != current_state) {
    physics.AddMassCmd(current_state);
    physics.AddSpringCmd(current_state, previous_state);
    if (win_conditions[current_preset](current_state)) {
      winning_states.insert(current_state);
    }
    visited_states.insert(current_state);
  }
}

auto StateManager::ClearGraph() -> void {
  winning_states.clear();
  visited_states.clear();
  physics.ClearCmd();
  physics.AddMassCmd(current_state);

  // The previous_state is no longer in the graph
  previous_state = current_state;

  // The starting state is no longer in the graph
  starting_state = current_state;
}

auto StateManager::FindWinningStates() -> void {
  winning_states.clear();
  std::unordered_map<State, int> state_masses;
  {
    std::lock_guard<LockableBase(std::mutex)> lock(physics.state.data_mtx);
    state_masses = physics.state.state_masses;
  }
  for (const auto &[state, mass] : state_masses) {
    if (win_conditions[current_preset](state)) {
      winning_states.insert(state);
    }
  }
}

auto StateManager::CurrentGenerator() -> StateGenerator {
  return generators[current_preset];
}

auto StateManager::CurrentWinCondition() -> WinCondition {
  return win_conditions[current_preset];
}