#!/usr/bin/python
#
# Copyright 2010 Google Inc. All Rights Reserved.

__author__ = 'asharif@google.com (Ahmad Sharif)'

from operator import attrgetter
import copy
import csv
import threading
import os.path

from automation.common import machine

DEFAULT_MACHINES_FILE = os.path.join(os.path.dirname(__file__), 'test_pool.csv')


class MachineManager(object):
  """Container for list of machines one can run jobs on."""

  @classmethod
  def FromMachineListFile(cls, filename):
    # Read the file and skip header
    csv_file = csv.reader(open(filename, 'rb'), delimiter=',', quotechar='"')
    csv_file.next()

    return cls([machine.Machine(hostname, label, cpu, int(cores), os, user)
                for hostname, label, cpu, cores, os, user in csv_file])

  def __init__(self, machines):
    self._machine_pool = machines
    self._lock = threading.RLock()

  def _GetMachine(self, mach_spec):
    available_pool = [m for m in self._machine_pool if mach_spec.IsMatch(m)]

    if available_pool:
      # find a machine with minimum uses
      uses = attrgetter('uses')

      mach = min(available_pool, key=uses)

      if mach_spec.preferred_machines:
        preferred_pool = [m for m in available_pool
                          if m.hostname in mach_spec.preferred_machines]
        if preferred_pool:
          mach = min(preferred_pool, key=uses)

      mach.Acquire(mach_spec.lock_required)

      return mach

  def GetMachines(self, required_machines):
    """Acquire machines for use by a job."""

    with self._lock:
      acquired_machines = [self._GetMachine(ms) for ms in required_machines]

      if not all(acquired_machines):
        # Roll back acquires
        while acquired_machines:
          mach = acquired_machines.pop()
          if mach:
            mach.Release()

      return acquired_machines

  def GetMachineList(self):
    with self._lock:
      return copy.deepcopy(self._machine_pool)

  def ReturnMachines(self, machines):
    with self._lock:
      for m in machines:
        m.Release()

  def __str__(self):
    return str(self._machine_pool)