# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging

from autotest_lib.client.common_lib import error
from autotest_lib.server.cros.faft.firmware_test import FirmwareTest


class firmware_ECUsbPorts(FirmwareTest):
    """
    Servo based EC USB port control test.
    """
    version = 1


    # Delay for remote shell command call to return
    RPC_DELAY = 1

    # Delay between turning off and on USB ports
    REBOOT_DELAY = 6

    # USB charge modes, copied from ec/include/usb_charge.h
    USB_CHARGE_MODE_DISABLED       = 0
    USB_CHARGE_MODE_SDP2           = 1
    USB_CHARGE_MODE_CDP            = 2
    USB_CHARGE_MODE_DCP_SHORT      = 3
    USB_CHARGE_MODE_ENABLED        = 4

    def initialize(self, host, cmdline_args):
        super(firmware_ECUsbPorts, self).initialize(host, cmdline_args)
        # Don't bother if there is no Chrome EC.
        if not self.check_ec_capability(['usb']):
            raise error.TestNAError("Nothing needs to be tested on this device")
        # Only run in normal mode
        self.switcher.setup_mode('normal')
        self.ec.send_command("chan 0")


    def cleanup(self):
        try:
            self.ec.send_command("chan 0xffffffff")
        except Exception as e:
            logging.error("Caught exception: %s", str(e))
        super(firmware_ECUsbPorts, self).cleanup()


    def fake_reboot_by_usb_mode_change(self):
        """
        Turn off USB ports and also kill FAFT client so that this acts like a
        reboot. If USB ports cannot be turned off or on, reboot step would
        fail.
        """
        for_all_ports_cmd = ('id=0; while [ $id -lt %d ];' +
                             'do ectool usbchargemode "$id" %d;' +
                             'id=$((id+1)); sleep 0.5; done')
        # Port disable - same for smart and dumb ports.
        ports_off_cmd = for_all_ports_cmd % (self._port_count,
                                             self.USB_CHARGE_MODE_DISABLED)
        # Port enable - different command based on smart/dumb port.
        port_enable_param = (self.USB_CHARGE_MODE_SDP2
            if self._smart_usb_charge else self.USB_CHARGE_MODE_ENABLED)
        ports_on_cmd = for_all_ports_cmd % (self._port_count, port_enable_param)
        cmd = ("sleep %d; %s; sleep %d; %s" %
               (self.RPC_DELAY, ports_off_cmd, self.REBOOT_DELAY,
                ports_on_cmd))
        block = False
        self.faft_client.system.run_shell_command(cmd, block)
        self.faft_client.disconnect()

    def __check_usb_enabled(self, idx):
        """Returns True if USB-A enable signal is high for a given index"""
        is_ioex = False
        gpio_name = 'USB%d_ENABLE' % (idx + 1)
        if self.faft_config.custom_usb_enable_pins:
            if idx >= len(self.faft_config.custom_usb_enable_pins):
                raise error.TestFail('No USB enable for index %d' % idx)
            is_ioex = self.faft_config.custom_usb_enable_pins[idx].get(
                    'ioex', False)
            gpio_name = self.faft_config.custom_usb_enable_pins[idx]['name']
        _, val = self.ec.send_command_get_output(
                '%sget %s' % (('gpio', 'ioex')[is_ioex], gpio_name),
                ['([01])[^\n\r]*\s%s' % gpio_name])[0]
        return val == '1'

    def get_port_count(self):
        """Get the number of USB ports."""
        for cnt in xrange(10):
            try:
                self.__check_usb_enabled(cnt)
            except error.TestFail:
                logging.info("Found %d USB ports", cnt)
                return cnt
        # Limit reached. Probably something went wrong.
        raise error.TestFail("Unexpected error while trying to determine " +
                             "number of USB ports")


    def check_power_off_mode(self):
        """Shutdown the system and check USB ports are disabled."""
        self.run_shutdown_cmd()
        self.wait_for('shutdown', 'Checking that all USB-A ports are disabled')
        # Check that all USB-A ports are disabled
        for idx in xrange(self._port_count):
            if self.__check_usb_enabled(idx):
                raise error.TestFail(
                        'Not all USB-A ports are disabled after shutdown')
        self.servo.power_short_press()


    def run_once(self):
        """Execute the main body of the test.
        """
        self._smart_usb_charge = (
            'smart_usb_charge' in self.faft_config.ec_capability)
        self._port_count = self.get_port_count()

        if self._port_count == 0:
            raise error.TestNAError("No USB-A port; nothing needs to be tested")

        if self.servo.main_device_is_ccd():
            logging.info("Using CCD, ignore checking USB port connection.")
        elif (self.servo.has_control('servo_v4_type') and
              self.servo.get('servo_v4_type') == 'type-c'):
            logging.info("Using type-c servo, ignore checking USB port connection.")
        elif (self.servo.has_control('servo_v4_type') and
              self.servo.get('servo_v4_type') != 'type-c'):
            # When only one USB-A port control is available, turning off the
            # USB-A port disconnects the network connection from the DUT.
            raise error.TestNAError("Only one USB-A port control; servo v4 type-C required")
        else:
            logging.info("Turn off all USB ports and then turn them on again.")
            self.switcher.mode_aware_reboot(
                    'custom', self.fake_reboot_by_usb_mode_change)

        logging.info("Check USB ports are disabled when powered off.")
        self.switcher.mode_aware_reboot('custom', self.check_power_off_mode)