#include "stdafx.h"
#include "L2Recruit.h"
#include "../Layer2Managers.h"
#include "../../../../PerceptualState/World/CurrentWorld.h"

static const size_t RECRUIT_TIMEOUT = 20;

L2Recruit::~L2Recruit()
{
  for (RecruitListType::iterator it = m_individualRecruits.begin(); it != m_individualRecruits.end(); ++it)
  {
    delete *it;
  }
}

void L2Recruit::onStartImplementation()
{
  L2TacticalAction::onStartImplementation();
  for (std::map<BWAPI::UnitType, size_t>::const_iterator it = m_units.begin();
       it != m_units.end();
       ++it)
  {
    for(size_t i = 0; i < it->second; ++i)
    {
      RecruitListID id = m_individualRecruits.insert(m_individualRecruits.end(), new L2IndividualRecruit(this, it->first));
      (*id)->execute(id);
    }
  }
}

void L2Recruit::onCancelImplementation()
{
  L2TacticalAction::onCancelImplementation();
  m_cancelled = true;
  RecruitListID it;
  RecruitListID nextIt = m_individualRecruits.begin();
  while(nextIt != m_individualRecruits.end())
  {
    it = nextIt;
    ++nextIt;
    (*it)->cancel(false);
  }
}

void L2Recruit::onEndImplementation()
{
  L2TacticalAction::onEndImplementation();
  m_cancelled = true;
  RecruitListID it;
  RecruitListID nextIt = m_individualRecruits.begin();
  while(nextIt != m_individualRecruits.end())
  {
    it = nextIt;
    ++nextIt;
    (*it)->terminate();
  }
}

void L2Recruit::onIndividualRecruitUnitTraining(const L2IndividualRecruit* recruit)
{
  if(!m_cancelled && m_isInfinite)
  {
    RecruitListID newId = m_individualRecruits.insert(m_individualRecruits.end(), new L2IndividualRecruit(this, recruit->getUnitToBuild()));
    (*newId)->execute(newId);
  }
}

void L2Recruit::onIndividualRecruitTerminated(RecruitListID id)
{
  Component::addZombieComponent(*id);
  m_individualRecruits.erase(id);
  if(m_individualRecruits.empty())
    terminate(Success);
}

L2Recruit::UnitRecordGroup L2Recruit::_getUnitsInvolved() const
{
  UnitRecordGroup unitsInvolved = Recruit::_getUnitsInvolved();   
  for(RecruitListType::const_iterator it = m_individualRecruits.begin(); it != m_individualRecruits.end(); ++it)
    (*it)->_getUnitsInvolved(unitsInvolved);
  return unitsInvolved;
}

// Class L2IndividualRecruit

void L2IndividualRecruit::_output(std::ostream& out) const
{
  out << "Individual recruit(" << m_unitToBuild << ")";
}

void L2IndividualRecruit::_getUnitsInvolved(L2Recruit::UnitRecordGroup& unitsInvolved) const
{
  if(m_buildingToUse != NULL)
  {
    unitsInvolved.insert(m_buildingToUse);
  }
}

void L2IndividualRecruit::execute(L2Recruit::RecruitListID id)
{  
  assert(m_unitToBuild.whatBuilds().first.isBuilding());
  m_id = id;

  if(CurrentWorld::getUnitTypeObject(m_unitToBuild)->isCreatable())
  {
    requestResources();
  } 
  else
  {
    // TODO: wait for CurrentWorld::getUnitTypeCreatableChangedEvent() instead
    m_unitCreatedSubscriberId = BWAPI::getUnitCreateEvent().subscribe(&m_unitCreatedEventHandler, NULL);
    m_currentState = WAITING_FOR_TECH;
  }
}

void L2IndividualRecruit::requestResources()
{
  Layer2Managers::getInstance().requestResources(
    m_l2Recruit->getPriority(), 
    ActionResources(Resources(m_unitToBuild),m_l2Recruit->getLocation(), m_unitToBuild.whatBuilds().first), 
    m_resourcesRequest);
  m_currentState = WAITING_FOR_RESOURCES;
}

void L2IndividualRecruit::onResourcesAllocated(void*, BWAPI::Unit* building)
{
  assert(building != NULL);
  assert(m_unitToBuild.whatBuilds().first == building->getType());
  assert(dynamic_cast<Layer2ManagersUnitData&>(BWAPI::getUnitData(building)).isInUse());  
  assert(building->isCompleted());
  assert(building->getTrainingQueue().empty());
  assert(!building->isUnpowered());

  m_buildingToUse = building;
  m_buildingDestroyedSubscriberId = BWAPI::getIndividualUnitDestroyEvent(m_buildingToUse).subscribe(&m_buildingDestroyedEventHandler,NULL);

  if(CurrentWorld::getUnitTypeObject(m_unitToBuild)->isCreatable())
  {
    Resources oldResources = Resources(BWAPI::Broodwar->self());
    assert(m_buildingToUse->canIssueCommand(BWAPI::UnitCommand::train(m_buildingToUse, m_unitToBuild)));
    bool actionSent = m_buildingToUse->train(m_unitToBuild);
    Layer2Managers::getInstance().terminateResourcesRequest(m_resourcesRequest);
    
    if(actionSent)
    {
      // TODO: assert crashes with BWAPI 3.7.4!
      //assert(oldResources == Resources(BWAPI::Broodwar->self()) + Resources(m_unitToBuild));
      m_l2Recruit->onIndividualRecruitUnitTraining(this);

      m_unitCreatedSubscriberId = BWAPI::getUnitCreateEvent().subscribe(&m_unitCreatedEventHandler, NULL);
    
      //TODO: getLatencyFrames and getRemainingLatencyFrames function do not return the same values, which one is the right one ?
      m_timeoutEvent.setWaitingTimeBeforeExecution(BWAPI::Broodwar->getLatencyFrames() + RECRUIT_TIMEOUT);
      m_timeoutEvent.activate();

      m_currentState = ACTION_SENT;
    }
    else
    {
      recruitFailed();
    }
  }
  else
  {
    recruitFailed();
  }
}

void L2IndividualRecruit::onUnitCreated(void*, BWAPI::Unit* unit)
{
  if(m_currentState == WAITING_FOR_TECH)
  {
    if(CurrentWorld::getUnitTypeObject(m_unitToBuild)->isCreatable())
    {
      BWAPI::getUnitCreateEvent().unsubscribe(m_unitCreatedSubscriberId);
      requestResources();
    }
  }
  else
  {
    if (m_buildingToUse == NULL || !m_buildingToUse->exists()) // Shouldn't happen?
    {
      recruitFailed();
    }
    else if(unit->getType() == m_unitToBuild && unit->getPosition() == m_buildingToUse->getPosition())
    {
      m_unitInConstruction = unit;
      BWAPI::getUnitCreateEvent().unsubscribe(m_unitCreatedSubscriberId);
      m_timeoutEvent.deactivate();

      // We don't register to onUnitCompleted now for better performance. -2 to make sure we don't miss the event.
      m_timeoutEvent.setWaitingTimeBeforeExecution(m_buildingToUse->getRemainingTrainTime()-2);
      m_timeoutEvent.activate();
      m_unitDestroyedSubscriberId = BWAPI::getIndividualUnitDestroyEvent(m_unitInConstruction).subscribe(&m_unitDestroyedEventHandler, NULL);
      
      m_currentState = UNIT_CREATED;

      if(m_cancelled)
      {
        m_buildingToUse->cancelTrain(0);
        recruitDone();
      }
    }
  } 
}

void L2IndividualRecruit::onUnitCompleted(void*, BWAPI::Unit* unit)
{
  if(m_unitInConstruction == unit)
  {
    m_l2Recruit->getUnitRecruitedEvent().raise(m_l2Recruit, m_unitInConstruction);
    recruitDone();
  }
}

void L2IndividualRecruit::onBuildingDestroyed(void*)
{
  recruitFailed();
}

void L2IndividualRecruit::onUnitDestroyed(void*)
{
  recruitFailed();
}

void L2IndividualRecruit::onTimeout()
{
  assert(m_currentState == ACTION_SENT || m_currentState == UNIT_CREATED);
  if(m_currentState == UNIT_CREATED)
  {
    m_unitCompletedSubscriberId = Spar->getPerceptualState()->getUnitCompletedEvent().subscribe(&m_unitCompletedEventHandler, NULL);
    m_timeoutEvent.deactivate();
    m_currentState = WAITING_TO_FINISH;
  } 
  else if(m_currentState == ACTION_SENT)
  {
    // This should never happen
    recruitFailed();
  }
}

void L2IndividualRecruit::recruitDone()
{
  if(m_buildingToUse != NULL)
  {
    dynamic_cast<Layer2ManagersUnitData&>(BWAPI::getUnitData(m_buildingToUse)).setInUseFlag(false);
    BWAPI::getIndividualUnitDestroyEvent(m_buildingToUse).unsubscribe(m_buildingDestroyedSubscriberId);
  }
  if(m_unitInConstruction != NULL)
  {
    BWAPI::getIndividualUnitDestroyEvent(m_unitInConstruction).unsubscribe(m_unitDestroyedSubscriberId);
  }
  Layer2Managers::getInstance().terminateResourcesRequest(m_resourcesRequest);
  BWAPI::getUnitCreateEvent().unsubscribe(m_unitCreatedSubscriberId);
  Spar->getPerceptualState()->getUnitCompletedEvent().unsubscribe(m_unitCompletedSubscriberId);
  m_timeoutEvent.deactivate();

  m_currentState = TERMINATED;
  m_l2Recruit->onIndividualRecruitTerminated(m_id);
}

void L2IndividualRecruit::recruitFailed()
{
  SPAR_LOG(LogTypes.LAYER_2_DECISION, "Individual Recruit Failed (%d) %s %d"
    , BWAPI::Broodwar->getFrameCount()
    , m_unitToBuild.getName().c_str()
    , (m_unitInConstruction==NULL ? -1 : m_unitInConstruction->getID()));
  if(m_buildingToUse != NULL)
  {
    dynamic_cast<Layer2ManagersUnitData&>(BWAPI::getUnitData(m_buildingToUse)).setInUseFlag(false);
    BWAPI::getIndividualUnitDestroyEvent(m_buildingToUse).unsubscribe(m_buildingDestroyedSubscriberId);
    m_buildingToUse = NULL;
  }
  if(m_unitInConstruction != NULL)
  {
    BWAPI::getIndividualUnitDestroyEvent(m_unitInConstruction).unsubscribe(m_unitDestroyedSubscriberId);
    m_unitInConstruction = NULL;
  }
  Layer2Managers::getInstance().terminateResourcesRequest(m_resourcesRequest);
  BWAPI::getUnitCreateEvent().unsubscribe(m_unitCreatedSubscriberId);
  Spar->getPerceptualState()->getUnitCompletedEvent().unsubscribe(m_unitCompletedSubscriberId);
  m_timeoutEvent.deactivate();

  if(m_cancelled)
  {
    m_currentState = TERMINATED;
    m_l2Recruit->onIndividualRecruitTerminated(m_id);
  }
  else
  {
    execute(m_id);
  }
}

void L2IndividualRecruit::cancel(bool cancelCurrentUnit)
{
  if(cancelCurrentUnit)
  {
    m_cancelled = true;
    switch(m_currentState)
    {
    case UNINITIALIZED:
      recruitDone();
      break;
    case WAITING_FOR_TECH:
      recruitDone();
      break;
    case WAITING_FOR_RESOURCES:
      recruitDone();
      break;
    case ACTION_SENT:
      // Wait until the unit is created (or the building is destroyed)
      break;
    case UNIT_CREATED:
    case WAITING_TO_FINISH:
      assert(m_buildingToUse != NULL && m_unitInConstruction != NULL);
      if(m_buildingToUse->exists() && !m_buildingToUse->getTrainingQueue().empty())
      {
        //TODO: Find the right unit in the queue to remove
        assert(m_buildingToUse->getTrainingQueue().size()==1);
        m_buildingToUse->cancelTrain(0);
      }
      recruitDone();
      break;
    case TERMINATED:
      break;
    }
  } 
  else 
  {
    switch(m_currentState)
    {
    case UNINITIALIZED:
      recruitDone();
      break;
    case WAITING_FOR_TECH:
      recruitDone();
      break;
    case WAITING_FOR_RESOURCES:
      recruitDone();
      break;
    case ACTION_SENT:
    case UNIT_CREATED:
    case WAITING_TO_FINISH:
    case TERMINATED:
      // Ignore the cancellation
      break;
    }
  }
}

void L2IndividualRecruit::terminate()
{
  m_cancelled = true;
  switch(m_currentState)
  {
  case UNINITIALIZED:
    recruitDone();
    break;
  case WAITING_FOR_TECH:
    recruitDone();
    break;
  case WAITING_FOR_RESOURCES:
    recruitDone();
    break;
  case ACTION_SENT:
    // At this point the unit will be lost, we cannot cancel it anymore.
    recruitDone();
    break;
  case UNIT_CREATED:
  case WAITING_TO_FINISH:
    assert(m_buildingToUse != NULL && m_unitInConstruction != NULL);
    if(m_buildingToUse->exists() && !m_buildingToUse->getTrainingQueue().empty())
    {
      //TODO: Find the right unit in the queue to remove
      assert(m_buildingToUse->getTrainingQueue().size()==1);
      m_buildingToUse->cancelTrain(0);
    }
    recruitDone();
    break;
  case TERMINATED:
    break;
  }
}