﻿/*
 * 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.
 */

namespace DotPulsar.Internal
{
    using System;
    using System.IO;
    using System.Net;
    using System.Net.Security;
    using System.Net.Sockets;
    using System.Security.Authentication;
    using System.Security.Cryptography.X509Certificates;
    using System.Threading.Tasks;

    public sealed class Connector
    {
        private readonly X509Certificate2Collection _clientCertificates;
        private readonly X509Certificate2? _trustedCertificateAuthority;
        private readonly bool _verifyCertificateAuthority;
        private readonly bool _verifyCertificateName;

        public Connector(
            X509Certificate2Collection clientCertificates,
            X509Certificate2? trustedCertificateAuthority,
            bool verifyCertificateAuthority,
            bool verifyCertificateName)
        {
            _clientCertificates = clientCertificates;
            _trustedCertificateAuthority = trustedCertificateAuthority;
            _verifyCertificateAuthority = verifyCertificateAuthority;
            _verifyCertificateName = verifyCertificateName;
        }

        public async Task<Stream> Connect(Uri serviceUrl)
        {
            var scheme = serviceUrl.Scheme;
            var host = serviceUrl.Host;
            var port = serviceUrl.Port;
            var encrypt = scheme == Constants.PulsarSslScheme;

            if (port == -1)
                port = encrypt ? Constants.DefaultPulsarSSLPort : Constants.DefaultPulsarPort;

            var stream = await GetStream(host, port).ConfigureAwait(false);

            if (encrypt)
                stream = await EncryptStream(stream, host).ConfigureAwait(false);

            return stream;
        }

        private static async Task<Stream> GetStream(string host, int port)
        {
            var tcpClient = new TcpClient();

            try
            {
                var type = Uri.CheckHostName(host);

                if (type == UriHostNameType.IPv4 || type == UriHostNameType.IPv6)
                    await tcpClient.ConnectAsync(IPAddress.Parse(host), port).ConfigureAwait(false);
                else
                    await tcpClient.ConnectAsync(host, port).ConfigureAwait(false);

                return tcpClient.GetStream();
            }
            catch
            {
                tcpClient.Dispose();
                throw;
            }
        }

        private async Task<Stream> EncryptStream(Stream stream, string host)
        {
            SslStream? sslStream = null;

            try
            {
                sslStream = new SslStream(stream, false, ValidateServerCertificate, null);
                await sslStream.AuthenticateAsClientAsync(host, _clientCertificates, SslProtocols.None, true).ConfigureAwait(false);
                return sslStream;
            }
            catch
            {
#if NETSTANDARD2_0
                if (sslStream is null)
                    stream.Dispose();
                else
                    sslStream.Dispose();
#else
                if (sslStream is null)
                    await stream.DisposeAsync().ConfigureAwait(false);
                else
                    await sslStream.DisposeAsync().ConfigureAwait(false);
#endif
                throw;
            }
        }

        private bool ValidateServerCertificate(object sender, X509Certificate? certificate, X509Chain? chain, SslPolicyErrors sslPolicyErrors)
        {
            if (sslPolicyErrors == SslPolicyErrors.None)
                return true;

            if (sslPolicyErrors.HasFlag(SslPolicyErrors.RemoteCertificateNotAvailable))
                return false;

            if (sslPolicyErrors.HasFlag(SslPolicyErrors.RemoteCertificateNameMismatch) && _verifyCertificateName)
                return false;

            if (sslPolicyErrors.HasFlag(SslPolicyErrors.RemoteCertificateChainErrors) && _verifyCertificateAuthority)
            {
                if (_trustedCertificateAuthority is null || chain is null || certificate is null)
                    return false;

                chain.ChainPolicy.ExtraStore.Add(_trustedCertificateAuthority);
                _ = chain.Build((X509Certificate2) certificate);

                for (var i = 0; i < chain.ChainElements.Count; i++)
                {
                    if (chain.ChainElements[i].Certificate.Thumbprint == _trustedCertificateAuthority.Thumbprint)
                        return true;
                }

                return false;
            }

            return true;
        }
    }
}