# Copyright 2021-2023 H2020 TeraFlow (https://www.teraflow-h2020.eu/)
#
# 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.

"""
P4 driver utilities.
"""

import logging
import queue
import sys
import threading
from functools import wraps
import grpc
import google.protobuf.text_format
from google.rpc import code_pb2

from p4.v1 import p4runtime_pb2
from p4.v1 import p4runtime_pb2_grpc

P4_ATTR_DEV_ID = 'id'
P4_ATTR_DEV_NAME = 'name'
P4_ATTR_DEV_VENDOR = 'vendor'
P4_ATTR_DEV_HW_VER = 'hw_ver'
P4_ATTR_DEV_SW_VER = 'sw_ver'
P4_ATTR_DEV_PIPECONF = 'pipeconf'

P4_VAL_DEF_VENDOR = 'Unknown'
P4_VAL_DEF_HW_VER = 'BMv2 simple_switch'
P4_VAL_DEF_SW_VER = 'Stratum'
P4_VAL_DEF_PIPECONF = 'org.onosproject.pipelines.fabric'

STREAM_ATTR_ARBITRATION = 'arbitration'
STREAM_ATTR_PACKET = 'packet'
STREAM_ATTR_DIGEST = 'digest'
STREAM_ATTR_UNKNOWN = 'unknown'

LOGGER = logging.getLogger(__name__)


class P4RuntimeException(Exception):
    """
    P4Runtime exception handler.

    Attributes
    ----------
    grpc_error : object
        gRPC error
    """

    def __init__(self, grpc_error):
        super().__init__()
        self.grpc_error = grpc_error

    def __str__(self):
        return str('P4Runtime RPC error (%s): %s',
                   self.grpc_error.code().name(), self.grpc_error.details())


def parse_p4runtime_error(fun):
    """
    Parse P4Runtime error.

    :param fun: function
    :return: parsed error
    """
    @wraps(fun)
    def handle(*args, **kwargs):
        try:
            return fun(*args, **kwargs)
        except grpc.RpcError as rpc_ex:
            raise P4RuntimeException(rpc_ex) from None
        except Exception as ex:
            raise Exception(ex) from None
    return handle


class P4RuntimeClient:
    """
    P4Runtime client.

    Attributes
    ----------
    device_id : int
        P4 device ID
    grpc_address : str
        IP address and port
    election_id : tuple
        Mastership election ID
    role_name : str
        Role name (optional)
    """
    def __init__(self, device_id, grpc_address, election_id, role_name=None):
        self.device_id = device_id
        self.election_id = election_id
        self.role_name = role_name
        self.stream_in_q = None
        self.stream_out_q = None
        self.stream = None
        self.stream_recv_thread = None
        LOGGER.debug(
            'Connecting to device %d at %s', device_id, grpc_address)
        self.channel = grpc.insecure_channel(grpc_address)
        self.stub = p4runtime_pb2_grpc.P4RuntimeStub(self.channel)
        try:
            self.set_up_stream()
        except P4RuntimeException:
            LOGGER.critical('Failed to connect to P4Runtime server')
            sys.exit(1)

    def set_up_stream(self):
        """
        Set up a gRPC stream.
        """
        self.stream_out_q = queue.Queue()
        # queues for different messages
        self.stream_in_q = {
            STREAM_ATTR_ARBITRATION: queue.Queue(),
            STREAM_ATTR_PACKET: queue.Queue(),
            STREAM_ATTR_DIGEST: queue.Queue(),
            STREAM_ATTR_UNKNOWN: queue.Queue(),
        }

        def stream_req_iterator():
            while True:
                st_p = self.stream_out_q.get()
                if st_p is None:
                    break
                yield st_p

        def stream_recv_wrapper(stream):
            @parse_p4runtime_error
            def stream_recv():
                for st_p in stream:
                    if st_p.HasField(STREAM_ATTR_ARBITRATION):
                        self.stream_in_q[STREAM_ATTR_ARBITRATION].put(st_p)
                    elif st_p.HasField(STREAM_ATTR_PACKET):
                        self.stream_in_q[STREAM_ATTR_PACKET].put(st_p)
                    elif st_p.HasField(STREAM_ATTR_DIGEST):
                        self.stream_in_q[STREAM_ATTR_DIGEST].put(st_p)
                    else:
                        self.stream_in_q[STREAM_ATTR_UNKNOWN].put(st_p)
            try:
                stream_recv()
            except P4RuntimeException as ex:
                LOGGER.critical('StreamChannel error, closing stream')
                LOGGER.critical(ex)
                for k in self.stream_in_q:
                    self.stream_in_q[k].put(None)
        self.stream = self.stub.StreamChannel(stream_req_iterator())
        self.stream_recv_thread = threading.Thread(
            target=stream_recv_wrapper, args=(self.stream,))
        self.stream_recv_thread.start()
        self.handshake()

    def handshake(self):
        """
        Handshake with gRPC server.
        """

        req = p4runtime_pb2.StreamMessageRequest()
        arbitration = req.arbitration
        arbitration.device_id = self.device_id
        election_id = arbitration.election_id
        election_id.high = self.election_id[0]
        election_id.low = self.election_id[1]
        if self.role_name is not None:
            arbitration.role.name = self.role_name
        self.stream_out_q.put(req)

        rep = self.get_stream_packet(STREAM_ATTR_ARBITRATION, timeout=2)
        if rep is None:
            LOGGER.critical('Failed to establish session with server')
            sys.exit(1)
        is_primary = (rep.arbitration.status.code == code_pb2.OK)
        LOGGER.debug('Session established, client is %s',
                        'primary' if is_primary else 'backup')
        if not is_primary:
            LOGGER.warning(
                'You are not the primary client, '
                'you only have read access to the server')

    def get_stream_packet(self, type_, timeout=1):
        """
        Get a new message from the stream.

        :param type_: stream type.
        :param timeout: time to wait.
        :return: message or None
        """
        if type_ not in self.stream_in_q:
            LOGGER.critical('Unknown stream type %s', type_)
            return None
        try:
            msg = self.stream_in_q[type_].get(timeout=timeout)
            return msg
        except queue.Empty:  # timeout expired
            return None

    @parse_p4runtime_error
    def get_p4info(self):
        """
        Retrieve P4Info content.

        :return: P4Info object.
        """

        LOGGER.debug('Retrieving P4Info file')
        req = p4runtime_pb2.GetForwardingPipelineConfigRequest()
        req.device_id = self.device_id
        req.response_type =\
            p4runtime_pb2.GetForwardingPipelineConfigRequest.P4INFO_AND_COOKIE
        rep = self.stub.GetForwardingPipelineConfig(req)
        return rep.config.p4info

    @parse_p4runtime_error
    def set_fwd_pipe_config(self, p4info_path, bin_path):
        """
        Configure the pipeline.

        :param p4info_path: path to the P4Info file
        :param bin_path: path to the binary file
        :return:
        """

        LOGGER.debug('Setting forwarding pipeline config')
        req = p4runtime_pb2.SetForwardingPipelineConfigRequest()
        req.device_id = self.device_id
        if self.role_name is not None:
            req.role = self.role_name
        election_id = req.election_id
        election_id.high = self.election_id[0]
        election_id.low = self.election_id[1]
        req.action =\
            p4runtime_pb2.SetForwardingPipelineConfigRequest.VERIFY_AND_COMMIT
        with open(p4info_path, 'r', encoding='utf8') as f_1:
            with open(bin_path, 'rb', encoding='utf8') as f_2:
                try:
                    google.protobuf.text_format.Merge(
                        f_1.read(), req.config.p4info)
                except google.protobuf.text_format.ParseError:
                    LOGGER.error('Error when parsing P4Info')
                    raise
                req.config.p4_device_config = f_2.read()
        return self.stub.SetForwardingPipelineConfig(req)

    def tear_down(self):
        """
        Tear connection with the gRPC server down.
        """

        if self.stream_out_q:
            LOGGER.debug('Cleaning up stream')
            self.stream_out_q.put(None)
        if self.stream_in_q:
            for k in self.stream_in_q:
                self.stream_in_q[k].put(None)
        if self.stream_recv_thread:
            self.stream_recv_thread.join()
        self.channel.close()
        del self.channel
