#pragma once
#include "../UnitRecord.h"
#include "UnitRecordType.h"
#include <set>

/**
 * @brief A transport unit record keeps track of what units are stored inside storage units. Storage units include
 *        Bunkers, Dropships, Shuttles and Overlords.
 */
class TransportUnitRecordType : public UnitRecordType
{
  TransportUnitRecordType(const TransportUnitRecordType &);
  TransportUnitRecordType & operator=(const TransportUnitRecordType &);
public:
  TransportUnitRecordType(UnitRecord* unitRecord, BWAPI::UnitType unitType)
    : UnitRecordType(unitRecord, unitType)
  {
    assert(unitType.spaceProvided() != 0);
  }

  class SavedState : public UnitRecordType::SavedState
  {
  public:
    SavedState(TransportUnitRecordType* unitRecordType);

    typedef std::set<AutoPtr<UnitRecord>>::const_iterator LoadedUnitIterator;
    std::pair<LoadedUnitIterator, LoadedUnitIterator> getLoadedUnits() const
    {
      return std::make_pair(m_loadedUnits.begin(), m_loadedUnits.end());
    }
  protected:
    std::set<AutoPtr<UnitRecord>> m_loadedUnits;
  };

  typedef std::set<AutoPtr<UnitRecord>>::const_iterator LoadedUnitIterator;
  std::pair<LoadedUnitIterator, LoadedUnitIterator> getLoadedUnits() const
  {
    return std::make_pair(m_loadedUnits.begin(), m_loadedUnits.end());
  }

  virtual void onDestroyed();

  virtual void onUnitDiscover();

  virtual void onUnitEvade();

  virtual SavedState* saveState() 
  {
    return new SavedState(this);
  }

protected:
  friend class UnitRecord;

  virtual void _outputUnitInfo(std::ostream& out) const;

  void addUnit(UnitRecord* unitRecord)
  {
    assert(m_loadedUnits.find(AutoPtr<UnitRecord>(*this, unitRecord)) == m_loadedUnits.end());
    m_loadedUnits.insert(AutoPtr<UnitRecord>(*this, unitRecord));
    m_loadedUnitsHistory.insert(AutoPtr<UnitRecord>(*this, unitRecord));
  }

  void removeUnit(UnitRecord* unitRecord)
  {
    m_loadedUnits.erase(AutoPtr<UnitRecord>(*this, unitRecord));
    m_loadedUnitsHistory.erase(AutoPtr<UnitRecord>(*this, unitRecord));
  }

  /**
   * @brief Units we are 100% sure are loaded inside this transport unit. Other units may
   *        be present but we do not have a clue what they can be.
   */
  std::set<AutoPtr<UnitRecord>> m_loadedUnits;

  /**
   * @brief History of units which were loaded in this transport unit but which we never witnessed the unloading.
   */
  std::set<AutoPtr<UnitRecord>> m_loadedUnitsHistory;
};
