/**
 * Copyright (c) 2015-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree. An additional grant
 * of patent rights can be found in the PATENTS file in the same directory.
 */

#pragma once

#include <cassert>
#include <cstdio>
#include <iostream>
#include <sstream>
#include <unordered_map>
#include <vector>

#include <torchcraft/primitives.h>
#include <torchcraft/unit.h>
#include <torchcraft/refcount.h>
#include "messages_generated.h"

#ifdef _MSC_VER
typedef unsigned int uint32_t;
typedef unsigned short uint16_t;
typedef int int32_t;
#endif

namespace torchcraft {
namespace replayer {

struct Resources;
struct Bullet;
struct Action;
class Frame;
class FrameDiff;
namespace detail { class UnitDiff; }

std::ostream& operator<<(std::ostream& out, const Frame& o);
std::istream& operator>>(std::istream& in, Frame& o);
std::ostream& operator<<(std::ostream& out, const FrameDiff& o);
std::istream& operator>>(std::istream& in, FrameDiff& o);

struct Bullet {
  int32_t type, x, y;
};

struct Action {
  std::vector<int32_t> action;
  int32_t uid;
  int32_t aid;
};

struct Resources {
  int32_t ore;
  int32_t gas;
  int32_t used_psi;
  int32_t total_psi;
  uint64_t upgrades;
  uint64_t upgrades_level;
  uint64_t techs;

  // clang-format off
  enum Upgrades : uint64_t {
    Terran_Infantry_Armor   = 1ll << 0,
    Terran_Vehicle_Plating  = 1ll << 1,
    Terran_Ship_Plating     = 1ll << 2,
    Zerg_Carapace           = 1ll << 3,
    Zerg_Flyer_Carapace     = 1ll << 4,
    Protoss_Ground_Armor    = 1ll << 5,
    Protoss_Air_Armor       = 1ll << 6,
    Terran_Infantry_Weapons = 1ll << 7,
    Terran_Vehicle_Weapons  = 1ll << 8,
    Terran_Ship_Weapons     = 1ll << 9,
    Zerg_Melee_Attacks      = 1ll << 10,
    Zerg_Missile_Attacks    = 1ll << 11,
    Zerg_Flyer_Attacks      = 1ll << 12,
    Protoss_Ground_Weapons  = 1ll << 13,
    Protoss_Air_Weapons     = 1ll << 14,
    Protoss_Plasma_Shields  = 1ll << 15,
    U_238_Shells            = 1ll << 16,
    Ion_Thrusters           = 1ll << 17,
    Titan_Reactor           = 1ll << 19,
    Ocular_Implants         = 1ll << 20,
    Moebius_Reactor         = 1ll << 21,
    Apollo_Reactor          = 1ll << 22,
    Colossus_Reactor        = 1ll << 23,
    Ventral_Sacs            = 1ll << 24,
    Antennae                = 1ll << 25,
    Pneumatized_Carapace    = 1ll << 26,
    Metabolic_Boost         = 1ll << 27,
    Adrenal_Glands          = 1ll << 28,
    Muscular_Augments       = 1ll << 29,
    Grooved_Spines          = 1ll << 30,
    Gamete_Meiosis          = 1ll << 31,
    Metasynaptic_Node       = 1ll << 32,
    Singularity_Charge      = 1ll << 33,
    Leg_Enhancements        = 1ll << 34,
    Scarab_Damage           = 1ll << 35,
    Reaver_Capacity         = 1ll << 36,
    Gravitic_Drive          = 1ll << 37,
    Sensor_Array            = 1ll << 38,
    Gravitic_Boosters       = 1ll << 39,
    Khaydarin_Amulet        = 1ll << 40,
    Apial_Sensors           = 1ll << 41,
    Gravitic_Thrusters      = 1ll << 42,
    Carrier_Capacity        = 1ll << 43,
    Khaydarin_Core          = 1ll << 44,
    Argus_Jewel             = 1ll << 47,
    Argus_Talisman          = 1ll << 49,
    Caduceus_Reactor        = 1ll << 51,
    Chitinous_Plating       = 1ll << 52,
    Anabolic_Synthesis      = 1ll << 53,
    Charon_Boosters         = 1ll << 54,
    Upgrade_60              = 1ll << 60,
    Unknow                  = 1ll << 62,
  };
  enum UpgradesLevel : uint64_t { // could be uint32_t
    Terran_Infantry_Armor_2   = 1ll << 0,
    Terran_Vehicle_Plating_2  = 1ll << 1,
    Terran_Ship_Plating_2     = 1ll << 2,
    Terran_Infantry_Weapons_2 = 1ll << 7,
    Terran_Vehicle_Weapons_2  = 1ll << 8,
    Terran_Ship_Weapons_2     = 1ll << 9,
    Zerg_Carapace_2           = 1ll << 3,
    Zerg_Flyer_Carapace_2     = 1ll << 4,
    Protoss_Ground_Armor_2    = 1ll << 5,
    Protoss_Air_Armor_2       = 1ll << 6,
    Zerg_Melee_Attacks_2      = 1ll << 10,
    Zerg_Missile_Attacks_2    = 1ll << 11,
    Zerg_Flyer_Attacks_2      = 1ll << 12,
    Protoss_Ground_Weapons_2  = 1ll << 13,
    Protoss_Air_Weapons_2     = 1ll << 14,
    Protoss_Plasma_Shields_2  = 1ll << 15,
    Terran_Infantry_Armor_3   = 1ll << 16,
    Terran_Vehicle_Plating_3  = 1ll << 17,
    Terran_Ship_Plating_3     = 1ll << 18,
    Terran_Infantry_Weapons_3 = 1ll << 23,
    Terran_Vehicle_Weapons_3  = 1ll << 24,
    Terran_Ship_Weapons_3     = 1ll << 25,
    Zerg_Carapace_3           = 1ll << 19,
    Zerg_Flyer_Carapace_3     = 1ll << 20,
    Protoss_Ground_Armor_3    = 1ll << 21,
    Protoss_Air_Armor_3       = 1ll << 22,
    Zerg_Melee_Attacks_3      = 1ll << 26,
    Zerg_Missile_Attacks_3    = 1ll << 27,
    Zerg_Flyer_Attacks_3      = 1ll << 28,
    Protoss_Ground_Weapons_3  = 1ll << 29,
    Protoss_Air_Weapons_3     = 1ll << 30,
    Protoss_Plasma_Shields_3  = 1ll << 31,
  };
  enum Techs : uint64_t {
    Stim_Packs         = 1ll << 0,
    Lockdown           = 1ll << 1,
    EMP_Shockwave      = 1ll << 2,
    Spider_Mines       = 1ll << 3,
    Scanner_Sweep      = 1ll << 4,
    Tank_Siege_Mode    = 1ll << 5,
    Defensive_Matrix   = 1ll << 6,
    Irradiate          = 1ll << 7,
    Yamato_Gun         = 1ll << 8,
    Cloaking_Field     = 1ll << 9,
    Personnel_Cloaking = 1ll << 10,
    Burrowing          = 1ll << 11,
    Infestation        = 1ll << 12,
    Spawn_Broodlings   = 1ll << 13,
    Dark_Swarm         = 1ll << 14,
    Plague             = 1ll << 15,
    Consume            = 1ll << 16,
    Ensnare            = 1ll << 17,
    Parasite           = 1ll << 18,
    Psionic_Storm      = 1ll << 19,
    Hallucination      = 1ll << 20,
    Recall             = 1ll << 21,
    Stasis_Field       = 1ll << 22,
    Archon_Warp        = 1ll << 23,
    Restoration        = 1ll << 24,
    Disruption_Web     = 1ll << 25,
    Unused_26          = 1ll << 26,
    Mind_Control       = 1ll << 27,
    Dark_Archon_Meld   = 1ll << 28,
    Feedback           = 1ll << 29,
    Optical_Flare      = 1ll << 30,
    Maelstrom          = 1ll << 31,
    Lurker_Aspect      = 1ll << 32,
    Unused_33          = 1ll << 33,
    Healing            = 1ll << 34,
    Nuclear_Strike     = 1ll << 45,
    Unknown            = 1ll << 46,
  };
  // clang-format on
};

class Frame : public RefCounted {
 public:
  // The keys of these hash tables are the players' ids.
  std::unordered_map<int32_t, std::vector<Unit>> units;
  std::unordered_map<int32_t, std::vector<Action>> actions;
  std::unordered_map<int32_t, Resources> resources;
  std::vector<Bullet> bullets;
  std::vector<uint8_t> creep_map; // Do not access directly
  uint32_t width, height;
  int reward;
  int is_terminal;

  Frame();
  Frame(Frame&& o);
  Frame(const Frame& o);
  Frame(const Frame* o);
  Frame& operator=(Frame other) noexcept;

  void swap(Frame& a, Frame& b);
  void clear();
  void filter(int32_t x, int32_t y, Frame& o) const;
  void combine(const Frame& next_frame);
  bool getCreepAt(uint32_t x, uint32_t y);
  
  flatbuffers::Offset<fbs::Frame> addToFlatBufferBuilder(flatbuffers::FlatBufferBuilder& builder) const;
  void readFromFlatBufferTable(const fbs::Frame& table);
}; // class Frame

// Frame diffs
class FrameDiff;

namespace detail {
  class UnitDiff {
   public:
    int id;
    std::vector<int32_t> var_ids;
    std::vector<int32_t> var_diffs;
    std::vector<int32_t> order_ids;
    std::vector<int32_t> order_diffs;
    int32_t order_size;
    double velocityX, velocityY;
    int64_t flags;
  };

  Frame* add(Frame* frame, FrameDiff* diff);
  void add(Frame* res, Frame* frame, FrameDiff* diff);

  inline bool orderUnitByiD(const Unit& a, const Unit& b) {
    return (a.id < b.id);
  }

  bool frameEq(Frame* f1, Frame* f2, bool debug = true);
} // namespace detail

class FrameDiff {
 public:
  std::vector<int32_t> pids;
  std::vector<std::vector<detail::UnitDiff>> units;
  // These are unlikely to be the same, so we just copy.
  std::unordered_map<int32_t, std::vector<Action>> actions;
  std::unordered_map<int32_t, Resources> resources;
  std::vector<Bullet> bullets;
  std::unordered_map<uint32_t, uint32_t> creep_map;
  // Width and height never changes, so we don't diff them
  int reward;
  int is_terminal;
  
  flatbuffers::Offset<fbs::FrameDiff> addToFlatBufferBuilder(flatbuffers::FlatBufferBuilder& builder) const;
  void readFromFlatBufferTable(const fbs::FrameDiff& fbsFrameDiff);
};

void writeTail(
    std::ostream& out,
    const std::unordered_map<int32_t, std::vector<replayer::Action>>& actions,
    const std::unordered_map<int32_t, replayer::Resources>& resources,
    const std::vector<replayer::Bullet>& bullets);

void readTail(
    std::istream& in,
    std::unordered_map<int32_t, std::vector<replayer::Action>>& actions,
    std::unordered_map<int32_t, replayer::Resources>& resources,
    std::vector<replayer::Bullet>& bullets);

// These diffing functions will order the IDs of units in each frame, and thus
// is not const.
FrameDiff frame_diff(Frame&, Frame&);
FrameDiff frame_diff(Frame*, Frame*);
Frame* frame_undiff(FrameDiff*, Frame*);
Frame* frame_undiff(Frame*, FrameDiff*);
void frame_undiff(Frame* result, FrameDiff*, Frame*);
void frame_undiff(Frame* result, Frame*, FrameDiff*);

} // namespace replayer
} // namespace torchcraft
