# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Fake test suite for local agent and real AMS integration test."""
import secrets
from typing import Any, Dict, List, Optional, Set, Tuple


_TEST_PROJECT = 'test-project'
_TEST_BRIGHTNESS_LEVEL = 150
_TEST_COLOR_HUE = 50
_TEST_COLOR_SATURATION = 30

# ======================== Supported Methods ========================== #
# See go/rainier-amsprocedure-api for more details

# Basic
_REBOOT = 'setReboot'
_FACTORY_RESET = 'setFactoryReset'

# On / Off
_TURN_ON_LIGHT = 'setOn'
_TURN_OFF_LIGHT = 'setOff'
_GET_LIGHT_STATE = 'getOnOff'

# Lock / Unlock
_LOCK_DEVICE = 'setLock'
_UNLOCK_DEVICE = 'setUnlock'
_GET_IS_LOCKED = 'getIsLocked'

# Brightness
_GET_LIGHT_BRIGHTNESS = 'getBrightness'
_SET_LIGHT_BRIGHTNESS = 'setBrightness'

# Color
_GET_LIGHT_COLOR = 'getColor'
_SET_LIGHT_COLOR = 'setColor'

# Test suite management
_START_SUITE = 'startTestSuite'
_END_SUITE = 'endTestSuite'
# ============================================================= #

# ====================== Helper methods ======================= #
# TODO(b/194029399) Retrieves correct test result id.
def generate_start_end_suite_rpc(
    device_ids: List[str]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Generates startTestSuite and endTestSuite RPC request.

    Args:
        device_ids: List of GDM device ids.

    Returns:
        Tuple: startTestSuite request and endTestSuite request.
    """
    fake_suite_id = secrets.token_hex(4)
    start_suite_rpc = {
        'projectId': _TEST_PROJECT,
        'method': _START_SUITE,
        'params': {'id': fake_suite_id,
                   'dutDeviceIds': device_ids}}
    end_suite_rpc = {
        'projectId': _TEST_PROJECT,
        'method': _END_SUITE,
        'params': {'id': fake_suite_id}}
    return start_suite_rpc, end_suite_rpc
# ============================================================= #


class BaseSuite:
    """Base class for all suite templates."""

    NAME = ''
    REQUIRED_TRAITS = set()

    def __init__(self, device_id: str):
        """Suite constructor.

        Args:
            device_id: GDM device id.
        """
        self._device_id = device_id
        self._procedures = []

    @classmethod
    def is_applicable_to(cls, device_traits: Set[str]) -> bool:
        """Whether the suite is applicable to the device traits.

        Args:
            device_traits: The device traits on HG.

        Returns:
            True if it's applicable, false otherwise.
        """
        return cls.REQUIRED_TRAITS.issubset(device_traits)

    def __str__(self) -> str:
        """Returns the suite name."""
        return self.NAME

    @property
    def device_id(self) -> str:
        """Returns the device id."""
        return self._device_id

    @property
    def procedures(self) -> List[Dict[str, Any]]:
        """Returns the procedures."""
        return self._procedures


class LightOnOffSuite(BaseSuite):
    """Light on off test suite."""

    NAME = 'Light OnOff Suite'
    REQUIRED_TRAITS = {'OnOff',}

    def __init__(self, device_id: str):
        """Light on off test suite constructor."""

        super().__init__(device_id=device_id)

        self._set_on_light_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _TURN_ON_LIGHT,
            'params': {'dutDeviceId': self._device_id}}

        self._set_off_light_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _TURN_OFF_LIGHT,
            'params': {'dutDeviceId': self._device_id}}

        self._get_state_light_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _GET_LIGHT_STATE,
            'params': {'dutDeviceId': self._device_id}}

        self._procedures = [
            self._set_on_light_rpc,
            self._get_state_light_rpc,
            self._set_off_light_rpc,
            self._get_state_light_rpc,
        ]


class LockUnlockSuite(BaseSuite):
    """Lock unlock test suite."""

    NAME = 'Lock Unlock Suite'
    REQUIRED_TRAITS = {'LockUnlock',}

    def __init__(self, device_id: str):
        """Lock unlock test suite constructor."""

        super().__init__(device_id=device_id)

        self._set_lock_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _LOCK_DEVICE,
            'params': {'dutDeviceId': self._device_id}}

        self._set_unlock_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _UNLOCK_DEVICE,
            'params': {'dutDeviceId': self._device_id}}

        self._get_locked_state_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _GET_IS_LOCKED,
            'params': {'dutDeviceId': self._device_id}}

        self._procedures = [
            self._set_lock_rpc,
            self._get_locked_state_rpc,
            self._set_unlock_rpc,
            self._get_locked_state_rpc,
        ]


class DeviceCommonSuite(BaseSuite):
    """Device common suite."""

    NAME = 'Device Common Suite'
    REQUIRED_TRAITS = {'Common',}

    def __init__(self, device_id: str):
        """Device common test suite constructor."""

        super().__init__(device_id=device_id)

        self._set_reboot_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _REBOOT,
            'params': {'dutDeviceId': self._device_id}}

        self._set_factory_reset_rpc = {
            'projectId': _TEST_PROJECT,
            'method': _FACTORY_RESET,
            'params': {'dutDeviceId': self._device_id}}

        self._procedures = [
            self._set_reboot_rpc,
            self._set_factory_reset_rpc,
        ]


class BrightnessSuite(BaseSuite):
    """Lighting brightness suite."""

    NAME = 'Brightness Suite'
    REQUIRED_TRAITS = {'OnOff',}

    def __init__(self, device_id: str):
        """Brightness test suite constructor."""

        super().__init__(device_id=device_id)

        self._set_brightness = {
            'projectId': _TEST_PROJECT,
            'method': _SET_LIGHT_BRIGHTNESS,
            'params': {'dutDeviceId': self._device_id,
                       'level':_TEST_BRIGHTNESS_LEVEL}}

        self._get_brightness = {
            'projectId': _TEST_PROJECT,
            'method': _GET_LIGHT_BRIGHTNESS,
            'params': {'dutDeviceId': self._device_id}}

        self._procedures = [
            self._set_brightness,
            self._get_brightness,
        ]


class ColorSuite(BaseSuite):
    """Lighting color suite."""

    NAME = 'Color Suite'
    REQUIRED_TRAITS = {'OnOff',}

    def __init__(self, device_id: str):
        """Color test suite constructor."""

        super().__init__(device_id=device_id)

        self._set_color = {
            'projectId': _TEST_PROJECT,
            'method': _SET_LIGHT_COLOR,
            'params': {'dutDeviceId': self._device_id,
                       'hue': _TEST_COLOR_HUE,
                       'saturation': _TEST_COLOR_SATURATION}}

        self._get_color = {
            'projectId': _TEST_PROJECT,
            'method': _GET_LIGHT_COLOR,
            'params': {'dutDeviceId': self._device_id}}

        self._procedures = [
            self._set_color,
            self._get_color,
        ]