#    -*- coding: utf-8 -*-
#
#     This file is part of sgx.py
#
#     Copyright (C) 2021 SKALE Labs
#
#     sgx.py is free software: you can redistribute it and/or modify
#     it under the terms of the GNU Affero General Public License as published
#     by the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     sgx.py is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU Affero General Public License for more details.
#
#     You should have received a copy of the GNU Affero General Public License
#     along with sgx.py.  If not, see <https://www.gnu.org/licenses/>.

import binascii
import json
import logging
import os
from dataclasses import dataclass
from enum import Enum
from time import sleep
from sgx.utils import public_key_to_address

from eth_utils.conversions import add_0x_prefix, remove_0x_prefix
import pem
import zmq
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, rsa, padding
from cryptography.hazmat.primitives.serialization import load_pem_private_key

from urllib.parse import urlparse

from sgx.constants import DEFAULT_TIMEOUT, SGX_ZMQ_RESPONSE_TIMEOUT_MS, CRT_FILENAME, KEY_FILENAME
from sgx.utils import SgxError


logger = logging.getLogger(__name__)


ZMQ_PORT = 1031
MAX_RETRIES = 10


class SgxZmqServerError(SgxError):
    pass


class SgxZmqUnreachableError(SgxError):
    pass


class DkgPolyStatus(Enum):
    FAIL = 0
    NEW_GENERATED = 1
    PREEXISTING = 2


@dataclass
class Account:
    name: str
    address: str
    public_key: str


@dataclass
class ComplaintResponse:
    share: str
    dh_key: str
    verification_vector_mult: str


class SgxZmq:
    def __init__(self, sgx_endpoint, path_to_cert=None, n=None, t=None):
        self.sgx_endpoint = get_provider(sgx_endpoint)
        self.path_to_cert = path_to_cert
        if n:
            self.n = n
        if t:
            self.t = t
        self.cert = self.__read_cert()
        self.ctx = zmq.Context()
        self.__init_method_types()
        self.sockets = dict()

    def __del__(self):
        for socket in self.sockets.values():
            socket.close()
        self.ctx.destroy()

    def ecdsa_sign(self, key_name, transaction_hash):
        params = dict()
        params['base'] = 10
        params['keyName'] = key_name
        params['messageHash'] = transaction_hash
        response = self.__send_request('ecdsaSignMessageHash', params)
        signature = response['result']
        vrs = (signature['signature_v'], signature['signature_r'], signature['signature_s'])
        return vrs

    def generate_key(self):
        params = dict()
        response = self.__send_request('generateECDSAKey', params)
        key_name = response['keyName']
        public_key = response['publicKey']
        public_key = add_0x_prefix(public_key)
        address = public_key_to_address(public_key)
        return Account(name=key_name, address=address, public_key=public_key)

    def get_public_key(self, keyName):
        params = dict()
        params['keyName'] = keyName
        response = self.__send_request('getPublicECDSAKey', params)
        publicKey = response['publicKey']
        return publicKey

    def generate_dkg_poly(self, poly_name):
        if self.is_poly_exists(poly_name):
            return DkgPolyStatus.PREEXISTING
        params = dict()
        params['polyName'] = poly_name
        params['t'] = self.t
        response = self.__send_request('generateDKGPoly', params)
        if response['status'] == 0:
            return DkgPolyStatus.NEW_GENERATED
        else:
            return DkgPolyStatus.FAIL

    def get_verification_vector(self, poly_name):
        params = dict()
        params['polyName'] = poly_name
        params['n'] = self.n
        params['t'] = self.t
        response = self.__send_request('getVerificationVector', params)
        verification_vector = response['verificationVector']
        return verification_vector

    def get_secret_key_contribution(self, poly_name, public_keys):
        public_keys = list(map(remove_0x_prefix, public_keys))
        params = dict()
        params['polyName'] = poly_name
        params['n'] = self.n
        params['t'] = self.t
        params['publicKeys'] = public_keys
        response = self.__send_request('getSecretShare', params)
        secret_key_contribution = response['secretShare']
        return secret_key_contribution

    def get_server_status(self):
        response = self.__send_request('getServerStatus')
        return response['status']

    def get_server_version(self):
        response = self.__send_request('getServerVersion')
        return response['version']

    def verify_secret_share(self, public_shares, eth_key_name, secret_share, index):
        params = dict()
        params['publicShares'] = public_shares
        params['ethKeyName'] = eth_key_name
        params['secretShare'] = secret_share
        params['n'] = self.n
        params['t'] = self.t
        params['index'] = index
        response = self.__send_request('dkgVerification', params)
        result = response['result']
        return result

    def create_bls_private_key(self, poly_name, bls_key_name, eth_key_name, secret_shares):
        params = dict()
        params['polyName'] = poly_name
        params['blsKeyName'] = bls_key_name
        params['ethKeyName'] = eth_key_name
        params['secretShare'] = secret_shares
        params['n'] = self.n
        params['t'] = self.t
        response = self.__send_request('createBLSPrivateKey', params)
        return response['status'] == 0

    def get_bls_public_key(self, bls_key_name):
        params = dict()
        params['blsKeyName'] = bls_key_name
        response = self.__send_request('getBLSPublicKeyShare', params)
        return response['blsPublicKeyShare']

    def complaint_response(self, poly_name, idx):
        params = dict()
        params['polyName'] = poly_name
        params['n'] = self.n
        params['t'] = self.t
        params['ind'] = idx
        response = self.__send_request('complaintResponse', params)
        return ComplaintResponse(
            share=response['share*G2'],
            dh_key=response['dhKey'],
            verification_vector_mult=response['verificationVectorMult'],
        )

    def mult_g2(self, to_mult):
        params = dict()
        params['x'] = to_mult
        response = self.__send_request('multG2', params)
        return response['x*G2']

    def import_bls_private_key(self, key_share_name, key_share):
        params = dict()
        params['keyShareName'] = key_share_name
        params['keyShare'] = key_share
        response = self.__send_request('importBLSKeyShare', params)
        encrypted_key = response['encryptedKeyShare']
        return encrypted_key

    def is_poly_exists(self, poly_name):
        params = dict()
        params['polyName'] = poly_name
        response = self.__send_request('isPolyExists', params)
        is_exists = response['IsExist']
        return is_exists

    def delete_bls_key(self, bls_key_name):
        params = dict()
        params['blsKeyName'] = bls_key_name
        response = self.__send_request('deleteBlsKey', params)
        result = response['deleted']

        return result

    def calculate_all_bls_public_keys(self, verification_vectors):
        params = dict()
        params['t'] = self.t
        params['n'] = self.n
        params['publicShares'] = verification_vectors
        response = self.__send_request('calculateAllBLSPublicKeys', params)
        result = response['publicKeys']

        return result

    def bls_sign(self, bls_key_name, message_hash):
        params = dict()
        params['keyShareName'] = bls_key_name
        params['messageHash'] = message_hash
        params['t'] = self.t
        params['n'] = self.n
        response = self.__send_request('blsSignMessageHash', params)
        result = response['signatureShare']

        return result

    def __send_request(self, method, params=None):
        if not params:
            params = dict()
        params['type'] = self.method_to_type[method]
        if self.path_to_cert:
            params['cert'] = self.cert
            msgSig = self.__sign_msg(params)
            params['msgSig'] = msgSig
        msg = json.dumps(params, separators=(',', ':'))
        p_id = os.getpid()
        if not self.sockets.get(p_id):
            socket_p_id = self.ctx.socket(zmq.DEALER)
            socket_p_id.setsockopt_string(zmq.IDENTITY, '135:14603077656239261618')
            socket_p_id.setsockopt(zmq.LINGER, 0)
            socket_p_id.setsockopt(zmq.SNDTIMEO, SGX_ZMQ_RESPONSE_TIMEOUT_MS)
            socket_p_id.setsockopt(zmq.RCVTIMEO, SGX_ZMQ_RESPONSE_TIMEOUT_MS)
            socket_p_id.connect(self.sgx_endpoint)
            self.sockets[p_id] = socket_p_id
        socket = self.sockets[p_id]
        socket.send_string(msg)
        # await reply
        response_str = None
        for _ in range(MAX_RETRIES):
            try:
                response_str = socket.recv().decode()
            except zmq.ZMQError:
                pass
            if response_str:
                break
            sleep(DEFAULT_TIMEOUT)
        if not response_str:
            raise SgxZmqUnreachableError('Max retries exceeded for sgx connection')
        response = json.loads(response_str)
        if (
            response.get('errorMessage') is not None and len(response.get('errorMessage'))
        ) or response['status']:
            raise SgxZmqServerError(response['errorMessage'])
        return response

    def __sign_msg(self, to_sign):
        msg = json.dumps(to_sign, separators=(',', ':'))
        msg = msg.replace(' ', '')
        key_path = os.path.join(self.path_to_cert, KEY_FILENAME)
        with open(key_path, 'rb') as key_file:
            private_key_bytes = key_file.read()
        private_key = load_pem_private_key(private_key_bytes, password=None)
        data = msg.encode()
        if isinstance(private_key, ec.EllipticCurvePrivateKey):
            signature = private_key.sign(data, ec.ECDSA(hashes.SHA256()))
        elif isinstance(private_key, rsa.RSAPrivateKey):
            signature = private_key.sign(data, padding.PKCS1v15(), hashes.SHA256())
        else:
            raise TypeError('Unsupported private key type for signing')
        return binascii.hexlify(signature).decode()

    def __read_cert(self):
        crt_path = os.path.join(self.path_to_cert, CRT_FILENAME)
        crt = pem.parse_file(crt_path)
        return str(crt[0])

    def __init_method_types(self):
        self.method_to_type = dict()
        self.method_to_type['ecdsaSignMessageHash'] = 'ECDSASignReq'
        self.method_to_type['generateECDSAKey'] = 'generateECDSAReq'
        self.method_to_type['getPublicECDSAKey'] = 'getPublicECDSAReq'
        self.method_to_type['generateDKGPoly'] = 'generateDKGPolyReq'
        self.method_to_type['getVerificationVector'] = 'getVerificationVectorReq'
        self.method_to_type['getSecretShare'] = 'getSecretShareReq'
        self.method_to_type['getServerStatus'] = 'getServerStatusReq'
        self.method_to_type['getServerVersion'] = 'getServerVersionReq'
        self.method_to_type['dkgVerification'] = 'dkgVerificationReq'
        self.method_to_type['createBLSPrivateKey'] = 'createBLSPrivateReq'
        self.method_to_type['getBLSPublicKeyShare'] = 'getBLSPublicReq'
        self.method_to_type['complaintResponse'] = 'complaintResponseReq'
        self.method_to_type['importBLSKeyShare'] = 'importBLSReq'
        self.method_to_type['isPolyExists'] = 'isPolyExistsReq'
        self.method_to_type['deleteBlsKey'] = 'deleteBLSKeyReq'
        self.method_to_type['calculateAllBLSPublicKeys'] = 'getAllBLSPublicReq'
        self.method_to_type['blsSignMessageHash'] = 'BLSSignReq'
        self.method_to_type['multG2'] = 'multG2Req'


def get_provider(endpoint):
    parsed_endpoint = urlparse(endpoint)
    ip, _ = parsed_endpoint.netloc.split(':')
    zmq_endpoint = 'tcp://' + ip + ':' + str(ZMQ_PORT)
    return zmq_endpoint
