using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using DateTime = System.DateTime;

namespace TcpRedirector
{
    internal class Program
    {
        private static int _proxyPort;

        public static string ProxyHost { get; private set; }

        public static int ProxyPort => _proxyPort;

        public static IReadOnlyList<byte> TargetHostBytes { get; private set; }
        public static IReadOnlyList<byte> TargetPortBytes { get; private set; }

        public static bool IsRunning { get; private set; }

        private static int Number = 0;

        public static int GetNumber()
        {
            return Interlocked.Increment(ref Number);
        }

        private static int Main(string[] args)
        {
            if (args.Length < 6)
            {
                Console.WriteLine("listenAddress listenPort proxyHost proxyPort targetHost targetPort");
                return -1;
            }

            if (!IPAddress.TryParse(args[0], out var listenAddress)) listenAddress = IPAddress.Any;
            if (!int.TryParse(args[1], out var listenPort)) listenPort = 0;
            ProxyHost = args[2];
            if (!int.TryParse(args[3], out _proxyPort)) _proxyPort = 0;
            var targetHost = args[4];
            if (!int.TryParse(args[5], out var targetPort)) targetPort = 0;

            TargetHostBytes = Encoding.ASCII.GetBytes(targetHost);
            TargetPortBytes = new[] { (byte)(targetPort << 8), (byte)targetPort };

            Console.Title = $"TR {listenAddress}:{listenPort} via {ProxyHost}:{ProxyPort} to {targetHost}:{targetPort}";

            IsRunning = true;
            var list = new List<Socks5Forwarder>();

            var startup = DateTime.Now;

            Task.Run(() =>
            {
                var listener = new TcpListener(new IPEndPoint(listenAddress, listenPort));
                listener.Start();

                Task<TcpClient> task = null;
                do
                {
                    if (task == null)
                    {
                        task = listener.AcceptTcpClientAsync();
                    }
                    else if (task.IsCompleted)
                    {
                        var client = task.Result;
                        lock (list)
                        {
                            list.Add(new Socks5Forwarder(client));
                        }
                        task = null;
                    }
                    Thread.Sleep(100);
                } while (IsRunning);

                listener.Stop();
            });

            while (IsRunning)
            {
                Console.Clear();

                Console.WriteLine($"UP {(DateTime.Now - startup).CompactFormat()} Since {startup}");

                //Print running info
                lock (list)
                {
                    Console.WriteLine();
                    Console.WriteLine($"Connections: {list.Count}");
                    Console.WriteLine();

                    foreach (var item in list)
                    {
                        Console.WriteLine(
                            $"#{item.Number:0000}" +
                            $" {item.CreatedTime:yyyy-MM-dd HH:mm:ss}"+
                            $" {item.IncomingEndPoint?.ToString() ?? "Unknown",20}" +
                            $"    {item.Status}" +
                            $"    Wai:{(DateTime.Now - item.LastActiveTime).CompactFormat(),6}" +
                            $"    Out:{item.OutgoingBytes.UnitFormat() + "B",10}" +
                            $"    In:{item.IncomingBytes.UnitFormat() + "B",10}" +
                            ""
                        );
                        if (string.IsNullOrWhiteSpace(item.Error) == false) Console.WriteLine($"\t{item.Error}");
                    }

                    //remove timeout items
                    foreach (var item in list.ToArray())
                    {
                        if ((DateTime.Now - item.LastActiveTime).TotalMinutes > 1)
                        {
                            if (item.Status
                                is Socks5ForwarderStatus.ERRBERU
                                or Socks5ForwarderStatus.ERRTOPR
                                or Socks5ForwarderStatus.ERRTOTA
                                or Socks5ForwarderStatus.DISCONN
                            )
                            {
                                list.Remove(item);
                            }
                        }
                    }
                }

                if (Console.KeyAvailable)
                {
                    switch (char.ToUpper(Console.ReadKey(true).KeyChar))
                    {
                        case 'Q':
                            IsRunning = false;
                            break;

                        default:
                            Console.WriteLine("Press Q to exit");
                            break;
                    }
                }

                Thread.Sleep(1000);
            }

            return 0;
        }
    }

    internal class Socks5Forwarder
    {
        public int Number { get;   }
        public Socks5ForwarderStatus Status { get; private set; } = Socks5ForwarderStatus.PENDING;

        public DateTime CreatedTime { get; } = DateTime.Now;

        public DateTime LastActiveTime { get; set; }

        public EndPoint? IncomingEndPoint { get; }
        public int IncomingBytes { get; private set; }
        public int OutgoingBytes { get; private set; }

        public string Error { get; private set; }

        public Socks5Forwarder(TcpClient incomingClient)
        {
            Number = Program.GetNumber();
            
            IncomingEndPoint = incomingClient.Client.RemoteEndPoint;
            
            Task.Run(() =>
            {
                Status = Socks5ForwarderStatus.CONTOPR;

                TcpClient outgoingClient = null;
                NetworkStream outgoingStream;

                try
                {
                    outgoingClient = new TcpClient(Program.ProxyHost, Program.ProxyPort);
                    outgoingStream = outgoingClient.GetStream();
                    LastActiveTime = DateTime.Now;

                    //socks 5 hand shake
                    outgoingStream.Write(new byte[] { 5, 1, 0 });
                    if (outgoingStream.ReadByte() != 5) throw new Exception("Server version error");
                    if (outgoingStream.ReadByte() != 0) throw new Exception("Server auth method error");
                }
                catch (Exception e)
                {
                    Status = Socks5ForwarderStatus.ERRTOPR;
                    Error = e.Message;
                    goto FINISHED;
                }

                Status = Socks5ForwarderStatus.CONTOTA;

                try
                {
                    LastActiveTime = DateTime.Now;
                    //socks 5 connection
                    var connect = new byte[] { 5, 1, 0, 3 }
                        .Concat(new[] { (byte)Program.TargetHostBytes.Count })
                        .Concat(Program.TargetHostBytes)
                        .Concat(Program.TargetPortBytes).ToArray();
                    outgoingStream.Write(connect);
                    if (outgoingStream.ReadByte() != 5) throw new Exception("Server version error");
                    var rep = outgoingStream.ReadByte();
                    if (rep != 0) throw new Exception($"Server response error:{rep}");
                    if (outgoingStream.ReadByte() != 0) throw new Exception($"Server response error");
                    var aty = outgoingStream.ReadByte();
                    var len = aty switch
                    {
                        1 => 4,
                        3 => outgoingStream.ReadByte(),
                        4 => 16,
                        _ => throw new Exception($"Server response AddressTYPE error")
                    };
                    for (var i = 0; i < len + 2; i++)
                    {
                        if (-1 == outgoingStream.ReadByte()) throw new Exception($"Server response padding error");
                    }
                }
                catch (Exception e)
                {
                    Status = Socks5ForwarderStatus.ERRTOTA;
                    Error = e.Message;
                    goto FINISHED;
                }

                NetworkStream incomingStream;
                try
                {
                    incomingClient.NoDelay = true;
                    outgoingClient.NoDelay = true;
                    incomingStream = incomingClient.GetStream();
                }
                catch (Exception e)
                {
                    Status = Socks5ForwarderStatus.ERRBERU;
                    Error = e.Message;
                    goto FINISHED;
                }

                Status = Socks5ForwarderStatus.RUNNING;
                LastActiveTime = DateTime.Now;

                //forward both

                var outgoingForward = Task.Run(() =>
                {
                    var buf = new byte[6000];
                    while (Program.IsRunning)
                    {
                        LastActiveTime = DateTime.Now;

                        int read;

                        try
                        {
                            read = incomingStream.Read(buf, 0, buf.Length);
                        }
                        catch (Exception e)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = $"Read Incoming: {e.Message}";
                            break;
                        }

                        if (read == 0)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = "Incoming Closed";
                            break;
                        }

                        OutgoingBytes += read;

                        try
                        {
                            outgoingStream.Write(buf, 0, read);
                        }
                        catch (Exception e)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = $"Write Outgoing:{e.Message}";
                            break;
                        }
                    }
                });

                var incomingForward = Task.Run(() =>
                {
                    var buf = new byte[6000];
                    while (Program.IsRunning)
                    {
                        LastActiveTime = DateTime.Now;

                        int read;

                        try
                        {
                            read = outgoingStream.Read(buf, 0, buf.Length);
                        }
                        catch (Exception e)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = $"Read Outgoing: {e.Message}";
                            break;
                        }

                        if (read == 0)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = "Outgoing disconnected";
                            break;
                        }

                        IncomingBytes += read;

                        try
                        {
                            incomingStream.Write(buf, 0, read);
                        }
                        catch (Exception e)
                        {
                            Status = Socks5ForwarderStatus.DISCONN;
                            Error = $"Write Incoming:{e.Message}";
                            break;
                        }
                    }
                });

                Task.WaitAny(outgoingForward, incomingForward);

                FINISHED:

                try
                {
                    outgoingClient?.Close();
                }
                catch
                {
                    //EAT ERROR
                }

                try
                {
                    incomingClient?.Close();
                }
                catch
                {
                    //EAT ERROR
                }
            });
        }
    }

    internal enum Socks5ForwarderStatus
    {
        PENDING,
        CONTOPR,
        CONTOTA,
        RUNNING,
        ERRTOPR,
        ERRTOTA,
        ERRBERU,
        DISCONN,
    }

    internal static class OutputFormatter
    {
        public static string CompactFormat(this TimeSpan value)
        {
            var lst = new List<string>();

            if (Math.Abs(value.TotalDays) >= 1) lst.Add($"{value.TotalDays:00}d");
            if (Math.Abs(value.Hours) >= 1) lst.Add($"{value.Hours:00}h");
            if (Math.Abs(value.Minutes) >= 1 && lst.Count < 2) lst.Add($"{value.Minutes:00}m");
            if (Math.Abs(value.Seconds) >= 1 && lst.Count < 2) lst.Add($"{value.Seconds:00}s");
            if (lst.Count < 2) lst.Add($"{value.Milliseconds:000}");

            return string.Join("", lst);
        }

        public static string UnitFormat(this int value, int digit = 3)
        {
            const string unit = " KMGTPE";

            var val = (double)value;

            var index = 0;
            while (val > 1000 && index < unit.Length)
            {
                val /= 1000;
                ++index;
            }

            return val.ToString("N" + digit) + unit[index];
        }
    }
}