//  
//   Rebex Sample Code License
// 
//   Copyright 2026, REBEX CR s.r.o.
//   All rights reserved.
//   https://www.rebex.net/
// 
//   Permission to use, copy, modify, and/or distribute this software for any
//   purpose with or without fee is hereby granted.
// 
//   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
//   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
//   OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
//   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
//   HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
//   WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
//   OTHER DEALINGS IN THE SOFTWARE.
// 

using Rebex.Net;
using System;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace Rebex.Samples
{
    /// <summary>
    /// Represents a tunnel between two endpoints.
    /// </summary>
    public class Tunnel
    {
        private const int BUFFER_SIZE = 64 * 1024; // size of buffers for receiving data

        private static int _nextId; // global ID counter

        private readonly object _sync; // synchronization object
        private readonly int _id; // ID of the tunnel
        private readonly TlsProxyConfiguration _config; // proxy configuration
        private readonly CancellationToken _cancellation; // cancellation token
        private readonly TunnelLogWriter _logger; // tunnel specific logger
        private readonly Action<int> _onClosed; // action to be called when the tunnel is closed

        private readonly TlsSocketDataProvider _incoming; // incoming connection
        private readonly TlsSocketDataProvider _outgoing; // outgoing connection
        private readonly EndPoint _incomingEndpoint; // IP address of the connected incoming connection
        private readonly EndPoint _outgoingEndpoint; // endpoint of the outgoing connection

        private bool _isClosed; // indicates whether the tunnel was closed

        /// <summary>
        /// Gets ID of the tunnel.
        /// </summary>
        public int Id { get { return _id; } }

        /// <summary>
        /// Determines whether the tunnel was stopped.
        /// </summary>
        private bool IsStopped { get { return _isClosed || _cancellation.IsCancellationRequested; } }

        /// <summary>
        /// Initializes new Tunnel object.
        /// </summary>
        public Tunnel(Socket incomingSocket, TlsProxyConfiguration config, CancellationToken cancellationToken, Action<int> onClosed)
        {
            if (incomingSocket == null) throw new ArgumentNullException("incomingSocket");
            if (config == null) throw new ArgumentNullException("config");
            if (onClosed == null) throw new ArgumentNullException("onClosed");

            // assign new ID
            _id = Interlocked.Increment(ref _nextId);

            // initialize private fields
            _sync = new object();
            _config = config;
            _logger = new TunnelLogWriter(config.LogWriter, Id);
            _cancellation = cancellationToken;
            _onClosed = onClosed;
            _incomingEndpoint = incomingSocket.RemoteEndPoint;
            _outgoingEndpoint = _config.TargetAddress;

            // initialize data providers
            _incoming = new TlsSocketDataProvider(new TlsServerSocket(incomingSocket), BUFFER_SIZE);
            _outgoing = new TlsSocketDataProvider(new TlsClientSocket(), BUFFER_SIZE);
        }

        /// <summary>
        /// Starts the tunnel.
        /// - Negotiate TLS layer at the incoming connection.
        /// - Connect to the outgoing endpoint.
        /// - Start data forwarding.
        /// </summary>
        public async Task Start()
        {
            LogWithTunnelDetails(LogLevel.Info, "Starting tunnel");

            await InitializeIncomingConnection((TlsServerSocket)_incoming.Socket).ConfigureAwait(false);
            await InitializeOutgoingConnection((TlsClientSocket)_outgoing.Socket).ConfigureAwait(false);

            LogWithTunnelDetails(LogLevel.Debug, "Tunnel established");

            // start forwarding data, no need to wait for task completion, only register error handler
            var dummy1 = StartForwarding(_incoming, _outgoing, "IN ==> OUT")
                                        .ContinueWith(TaskHelpers.GetTaskLogErrorFunc(Log, "Forwarding task failed (IN ==> OUT)"),
                                                      TaskContinuationOptions.OnlyOnFaulted);
            var dummy2 = StartForwarding(_outgoing, _incoming, "IN <== OUT")
                                        .ContinueWith(TaskHelpers.GetTaskLogErrorFunc(Log, "Forwarding task failed (IN <== OUT)"),
                                                      TaskContinuationOptions.OnlyOnFaulted);
        }

        /// <summary>
        /// Initializes incoming connection.
        /// </summary>
        private async Task InitializeIncomingConnection(TlsServerSocket socket)
        {
            socket.LogWriter = _logger.GetSpecific("IN"); // assign specific logger
            socket.Timeout = _config.Timeout; // set desired timeout
            socket.Parameters.Certificate = _config.Certificate; // set proxy/server certificate
            socket.Parameters.Version = _config.TlsVersions; // set desired TLS versions
            socket.ClientHelloReceived += (s, e) => // register event to display SNI info
            {
                Log(LogLevel.Debug, "Received SNI: {0}.", e.ServerName);
            };

            // perform TLS negotiation
            await socket.NegotiateAsync().ConfigureAwait(false);
        }

        /// <summary>
        /// Initializes outgoing connection.
        /// </summary>
        private async Task InitializeOutgoingConnection(TlsClientSocket socket)
        {
            socket.LogWriter = _logger.GetSpecific("OUT"); // assign specific logger
            socket.Timeout = _config.Timeout; // set desired timeout

            // connect to desired outgoing endpoint
            await socket.ConnectAsync(_config.TargetAddress).ConfigureAwait(false);
        }

        /// <summary>
        /// Starts data forwarding and keeps the loop running while connections are active or time out is received.
        /// </summary>
        private async Task StartForwarding(TlsSocketDataProvider reader, TlsSocketDataProvider writer, string direction)
        {
            string operation = "forwarding";
            try
            {
                while (!IsStopped)
                {					
                    operation = "receiving"; // wait for data
                    int readCount = await reader.ReceiveAsync().ConfigureAwait(false);
                    if (readCount == 0 || IsStopped)
                        return;

                    Log(LogLevel.Debug, "Forwarding {0} bytes ({1}).", readCount, direction);
                                        
                    operation = "sending"; // forward received data
                    await writer.SendAsync(ArraySegmentHelper.Slice(reader.Buffer, 0, readCount)).ConfigureAwait(false);
                }
            }
            catch (Exception ex)
            {
                // do not log timeout error, it is expected
                Exception x = ex;
                while (x != null)
                {
                    if (x is TimeoutException)
                        return;

                    var tx = x as TlsException;
                    if (tx != null && tx.Status == NetworkSessionExceptionStatus.Timeout)
                        return;

                    x = x.InnerException;
                }
                Log(ex, string.Format("Error while {0} data ({1})", operation, direction));
            }
            finally
            {
                await CloseAsync().ConfigureAwait(false);
            }
        }

        /// <summary>
        /// Closes the tunnel.
        /// </summary>
        public async Task CloseAsync()
        {
            lock (_sync)
            {
                if (_isClosed)
                    return;
                _isClosed = true;
            }

            try
            {
                LogWithTunnelDetails(LogLevel.Info, "Closing tunnel");

                try
                {
                    await _incoming.CloseAsync().ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    Log(LogLevel.Debug, "Error while closing incoming tunnel: {0}", ex);
                }

                try
                {
                    await _outgoing.CloseAsync().ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    Log(LogLevel.Debug, "Error while closing outgoing tunnel: {0}", ex);
                }

                LogWithTunnelDetails(LogLevel.Info, "Tunnel closed");
            }
            finally
            {
                _onClosed(Id);
            }
        }

        /// <summary>
        /// Helper logging method - logs error.
        /// </summary>
        private void Log(Exception ex, string message)
        {
            if (_logger.Level != LogLevel.Off && !IsStopped)
            {
                Log(LogLevel.Error, "{0}: {1}", message, ex);
            }
        }

        /// <summary>
        /// Helper logging method - logs message with Tunnel endpoints info.
        /// </summary>
        private void LogWithTunnelDetails(LogLevel level, string message)
        {
            if (_logger.Level != LogLevel.Off)
            {
                Log(level, $"{message} ({_incomingEndpoint}) --TLS--> ({_config.ListeningAddress.Port}) --plain--> ({_outgoingEndpoint}).");
            }
        }

        /// <summary>
        /// Helper logging method.
        /// </summary>
        private void Log(LogLevel level, string format, params object[] args)
        {
            if (_logger.Level != LogLevel.Off)
            {
                Log(level, string.Format(format, args));
            }
        }

        /// <summary>
        /// Helper logging method.
        /// </summary>
        private void Log(LogLevel level, string message)
        {
            if (_logger.Level != LogLevel.Off)
            {
                _logger.Write(level, typeof(Tunnel), Id, "INFO", message);
            }
        }
    }
}
