from nullptr.missions.base import BaseMission, MissionParam
from nullptr.models.waypoint import Waypoint
from nullptr.models.survey import Survey
from nullptr.models.contract import Contract
from nullptr.util import *
from nullptr.missions.extraction import ExtractionMission

class MiningMission(ExtractionMission): 
  @classmethod
  def params(cls):
    return {
      'site': MissionParam(Waypoint, True),
      'resources': MissionParam(list, True)
    }
  
  def start_state(self):
    return 'travel-to'
    
  def steps(self):
    return {
      **self.travel_steps('to', 'site', 'extract'),
        'extract': (self.step_extract, {
            'more': 'extract',
            'done': 'unload'
        }),
        'unload': (self.step_unload, {
            'more': 'unload',
            'done': 'done'
        })
    }
  
  def get_survey(self):
    resources = self.st('resources')
    site = self.rst(Waypoint,'site')
    best_score = 0
    best_survey = None
    # todo optimize
    for s in self.store.all(Survey):
      if site != s.waypoint:
        continue
      good = len([1 for r in s.deposits if r in resources])
      total = len(s.deposits)
      score = good / total
      if score > best_score:
        best_score = score
        best_survey = s
    return best_survey
    
  def step_extract(self):
    survey = self.get_survey()
    result = self.api.extract(self.ship, survey)
    symbol = sg(result,'extraction.yield.symbol')
    units  = sg(result,'extraction.yield.units')
    self.next_step = self.ship.cooldown
    if self.ship.cargo_space() > 5:
      return 'more'
    else:
      return 'done'