import dataclasses
import logging

from siotls.configuration import TLSNegotiatedConfiguration
from siotls.contents import ChangeCipherSpec, alerts
from siotls.contents.handshakes import (
    CertificateHandshake,
    CertificateRequest,
    CertificateVerify,
    EncryptedExtensions,
    Finished,
    HelloRetryRequest,
    ServerHello,
)
from siotls.contents.handshakes.certificate import RawPublicKeyEntry, X509Entry
from siotls.contents.handshakes.extensions import (
    ALPN,
    ClientCertificateTypeResponse,
    KeyShareResponse,
    KeyShareRetry,
    ServerCertificateTypeResponse,
    SignatureAlgorithms,
    SupportedVersionsResponse,
)
from siotls.crypto import TLSCipherSuite, TLSKeyExchange, TLSSignatureScheme
from siotls.iana import (
    CertificateType,
    ContentType,
    ExtensionType,
    HandshakeType,
    TLSVersion,
)

from .. import CERTIFICATE_VERIFY_SERVER, State
from . import ServerWaitCertificate, ServerWaitFinished

logger = logging.getLogger(__name__)


class ServerWaitClientHello(State):
    can_send = True
    can_send_application_data = False

    def __init__(self, connection):
        super().__init__(connection)
        self._is_first_client_hello = True

    def process(self, client_hello):
        if (client_hello.content_type != ContentType.HANDSHAKE
            or client_hello.msg_type != HandshakeType.CLIENT_HELLO):
            super().process(client_hello)
            return

        if not self._is_first_client_hello:
            self._find_common_cipher_suite(client_hello.cipher_suites)
            if client_hello.random != self._client_unique:
                e = "client's random cannot change in between Hellos"
                raise alerts.IllegalParameter(e)

        if self._is_first_client_hello:
            self._setup(client_hello)

        server_exts, shared_key = self._negociate(client_hello.extensions)
        if not shared_key:
            self._send_hello_retry_request(client_hello.legacy_session_id, server_exts)
            return
        self._send_server_hello(client_hello.legacy_session_id, server_exts, shared_key)

        if self.config.require_peer_authentication:
            self._request_user_certificate()

        self._send_server_certificate()

        self._send_finished()
        server_finished_th = self._transcript.digest()
        if self.config.require_peer_authentication:
            self._move_to_state(ServerWaitCertificate, server_finished_th)
        else:
            self._move_to_state(
                ServerWaitFinished, server_finished_th, server_finished_th)

    def _setup(self, client_hello):
        self.nconfig = TLSNegotiatedConfiguration()
        self.nconfig.cipher_suite = self._find_common_cipher_suite(
            client_hello.cipher_suites)
        self._client_unique = client_hello.random
        self._cipher = TLSCipherSuite[self.nconfig.cipher_suite](
            'server', self._client_unique, log_keys=self.config.log_keys)
        self._transcript.post_init(self._cipher.digestmod)

    def _find_common_cipher_suite(self, cipher_suites):
        for cipher_suite in self.config.cipher_suites:
            if cipher_suite in cipher_suites:
                return cipher_suite
        e = "no common cipher suite found"
        raise alerts.HandshakeFailure(e)

    def _send_hello_retry_request(self, session_id, server_extensions):
        if not self._is_first_client_hello:
            e = "invalid KeyShare in second ClientHello"
            raise alerts.HandshakeFailure(e)
        self._is_first_client_hello = False

        # make sure the client doesn't change its algorithms in between flights
        self.config = dataclasses.replace(self.config,
            cipher_suites=[self._cipher.iana_id],
            key_exchanges=[self.nconfig.key_exchange],
            signature_algorithms=[self._signature.iana_id]
        )

        clear_extensions, _ = server_extensions
        self._transcript.do_hrr_dance()
        self._send_content(HelloRetryRequest(
            HelloRetryRequest.random,
            session_id,
            self._cipher.iana_id,
            clear_extensions,
        ))
        self._send_content(ChangeCipherSpec())

    def _send_server_hello(self, session_id, extensions, shared_key):
        clear_extensions, encrypted_extensions = extensions

        self._cipher.skip_early_secrets()
        self._send_content(ServerHello(
            self._server_unique,
            session_id,
            self._cipher.iana_id,
            clear_extensions,
        ))
        if self._is_first_client_hello:
            self._send_content(ChangeCipherSpec())
        self._cipher.derive_handshake_secrets(shared_key, self._transcript.digest())
        self._send_content(EncryptedExtensions(encrypted_extensions))

    def _request_user_certificate(self):
        # we ignore POST_HANDSHAKE_AUTH at the moment
        certificate_request_extensions = [
            SignatureAlgorithms(self.config.signature_algorithms),
        ]
        self._send_content(CertificateRequest(
            certificate_request_context=b'',  # only for POST_HANDSHAKE_AUTH
            extensions=certificate_request_extensions,
        ))

    def _send_server_certificate(self):
        certificate_list = []
        match self.nconfig.server_certificate_type:
            case CertificateType.RAW_PUBLIC_KEY:
                certificate_list.append(RawPublicKeyEntry(self.config.public_key, []))
            case CertificateType.X509:
                certificate_list.extend([
                    X509Entry(certificate, [])
                    for certificate
                    in self.config.certificate_chain
                ])
            case unknown:
                e = f"unknown server certificate type: {unknown!r}"
                raise NotImplementedError(e)

        self._send_content(CertificateHandshake(b'', certificate_list))
        self._send_content(CertificateVerify(
            self._signature.iana_id,
            self._signature.sign(
                CERTIFICATE_VERIFY_SERVER + self._transcript.digest(),
            )
        ))

    def _send_finished(self):
        self._send_content(Finished(
            self._cipher.sign_finish(self._transcript.digest())
        ))

    def _negociate(self, client_extensions):
        clear_extensions = []
        encrypted_extensions = []

        def negociate(ext_name, *args, **kwargs):
            ext_type = getattr(ExtensionType, ext_name.upper())
            ext = client_extensions.get(ext_type)
            meth = getattr(self, f'_negociate_{ext_name}')
            clear_exts, encrypted_exts, *rest = meth(ext, *args, **kwargs)
            clear_extensions.extend(clear_exts)
            encrypted_extensions.extend(encrypted_exts)
            return rest[0] if rest else None

        negociate('supported_versions')
        negociate('supported_groups')
        negociate('server_certificate_type')
        negociate('signature_algorithms')

        shared_key = negociate('key_share')
        if not shared_key:
            return (clear_extensions, encrypted_extensions), None

        negociate('client_certificate_type')

        negociate('max_fragment_length')
        negociate('application_layer_protocol_negotiation')

        return (clear_extensions, encrypted_extensions), shared_key

    def _negociate_supported_versions(self, supported_versions_ext):
        if not supported_versions_ext:
            e = "client doesn't support TLS 1.3"
            raise alerts.ProtocolVersion(e)
        if TLSVersion.TLS_1_3 not in supported_versions_ext.versions:
            e = "client doesn't support TLS 1.3"
            raise alerts.ProtocolVersion(e)
        return [SupportedVersionsResponse(TLSVersion.TLS_1_3)], []

    def _negociate_supported_groups(self, supported_groups_ext):
        if not supported_groups_ext:
            raise alerts.MissingExtension(ExtensionType.SUPPORTED_GROUPS)
        for group in self.config.key_exchanges:
            if group in supported_groups_ext.named_group_list:
                self.nconfig.key_exchange = group
                return [], []
        e = "no common key exchange found"
        raise alerts.HandshakeFailure(e)

    def _negociate_client_certificate_type(self, cct_ext):
        if not self.config.require_peer_authentication:
            return [], []
        client_cert_types = (
            cct_ext.certificate_types if cct_ext else [CertificateType.X509]
        )
        for cert_type in self.config.peer_certificate_types:
            if cert_type in client_cert_types:
                self.nconfig.client_certificate_type = cert_type
                if cct_ext:
                    return [], [ClientCertificateTypeResponse(cert_type)]
                return [], []
        e = "no common client certificate type found"
        raise alerts.UnsupportedCertificate(e)

    def _negociate_server_certificate_type(self, sct_ext):
        server_cert_types = (
            sct_ext.certificate_types if sct_ext else [CertificateType.X509]
        )
        for cert_type in self.config.certificate_types:
            if cert_type in server_cert_types:
                self.nconfig.server_certificate_type = cert_type
                if sct_ext:
                    return [], [ServerCertificateTypeResponse(cert_type)]
                return [], []
        e = "no common server certificate type found"
        raise alerts.UnsupportedCertificate(e)

    def _negociate_signature_algorithms(self, sa_ext):
        if not sa_ext:
            raise alerts.MissingExtension(ExtensionType.SIGNATURE_ALGORITHMS)
        sign_algo_ids = {
            suite.iana_id
            for suite
            in TLSSignatureScheme.for_key_algo(self.config.asn1_public_key['algorithm'])
            if suite.iana_id in self.config.private_key_signature_algorithms
        }
        for sign_algo_id in sa_ext.supported_signature_algorithms:
            if sign_algo_id in sign_algo_ids:
                SignScheme = TLSSignatureScheme[sign_algo_id]
                self._signature = SignScheme(private_key=self.config.private_key)
                self.nconfig.signature_algorithm = sign_algo_id
                return [], []
        e = "no common signature algorithm found"
        raise alerts.HandshakeFailure(e)

    def _negociate_key_share(self, key_share_ext):
        key_exchange = self.nconfig.key_exchange
        if key_share_ext and key_exchange in key_share_ext.client_shares:
            # possible to resume key share => ServerHello
            client_exchange = key_share_ext.client_shares[key_exchange]
            KeyExchange = TLSKeyExchange[key_exchange]
            try:
                private_key, server_exchange = KeyExchange.init()
                shared_key = KeyExchange.resume(private_key, client_exchange)
            except ValueError as exc:
                e = "error while resuming key share"
                raise alerts.HandshakeFailure(e) from exc
            response = KeyShareResponse(key_exchange, server_exchange)
        else:
            # impossible to resume key share => HelloRetryRequest
            shared_key = None
            response = KeyShareRetry(key_exchange)
        return [response], [], shared_key

    def _negociate_max_fragment_length(self, mfl_ext):
        if mfl_ext:
            self.nconfig.max_fragment_length = mfl_ext.octets
            return [], [mfl_ext]

        self.nconfig.max_fragment_length = self.config.max_fragment_length
        return [], []

    def _negociate_application_layer_protocol_negotiation(self, alpn_ext):
        if not alpn_ext or not self.config.alpn:
            self.nconfig.alpn = None
            return [], []
        for proto in self.config.alpn:
            if proto in alpn_ext.protocol_name_list:
                self.nconfig.alpn = proto
                return [], [ALPN([proto])]
        e = "no common application layer protocol found"
        raise alerts.NoApplicationProtocol(e)
    _negociate_alpn = _negociate_application_layer_protocol_negotiation
