using System;
using System.Collections.Generic;
using System.Threading;
using System.IO;
using Bitvise;
using Bitvise.FlowSshNet;

namespace FlowSshNet_ExecStress
{
    class ExecCommand
    {
        public uint m_cmdRepeatCnt;
        public string m_cmd;
    }

    class Coordination
    {
        public ManualResetEvent m_stopEvent = new ManualResetEvent(false);
        public string m_stopReason;

        public void Stop(string reason)
        {
            lock (this)
            {
                if (m_stopReason == null)
                    m_stopReason = reason;

                m_stopEvent.Set();
            }
        }
    }

    class ExecErr : System.Exception
    {
        public ExecErr(string msg) : base(msg) { }
    }

    class ExecThread
    {
        public Coordination m_coord;
        public uint         m_threadNr;

        public string       m_host;
        public uint         m_port;
        public string       m_hkfp;
        public string       m_user;
        public string       m_pass;
        public ExecCommand  m_cmd;

        public void Run()
        {
            try
            {
                Console.WriteLine("Thread {0} started", m_threadNr);

                using (System.IO.StreamWriter file = new System.IO.StreamWriter(@".\DataRcvdThread_" + m_threadNr + ".txt"))
                {
                    for (uint loop = 1; loop <= m_cmd.m_cmdRepeatCnt; ++loop)
                    {
                        using (Client client = new Client())
                        {
                            client.SetAppName("FlowSshNet_SftpStress");

                            client.SetHost(m_host);
                            client.SetPort(m_port);
                            client.SetUserName(m_user);
                            client.SetPassword(m_pass);
                            client.AcceptHostKeySha256(m_hkfp);

                            Console.WriteLine("Thread {0} connecting for {1}th time", m_threadNr, loop);

                            ProgressHandler ph = new ProgressHandler();
                            client.Connect(ph);
                            ph.WaitDone();
                            if (!ph.Success)
                                throw new ExecErr(ph.DescribeConnectError());

                            using (ClientSessionChannel channel = new ClientSessionChannel(client))
                            {
                                ph = new ProgressHandler();
                                channel.OpenRequest(ph);
                                ph.WaitDone();
                                if (!ph.Success)
                                {
                                    string auxInfo = (ph.GetAuxInfo().Length > 0) ? ph.GetAuxInfo() : "(no additional info)";
                                    throw new ExecErr("Opening session channel failed: " + auxInfo);
                                }

                                ph = new ProgressHandler();
                                channel.ExecRequest(m_cmd.m_cmd, ph);
                                ph.WaitDone();
                                if (!ph.Success)
                                {
                                    string auxInfo = (ph.GetAuxInfo().Length > 0) ? ph.GetAuxInfo() : "(no additional info)";
                                    throw new ExecErr("Session channel - failed to execute remote command: " + auxInfo);
                                }

                                Int32 receiveCallCnt = 0;
                                Int32 totalBytesReceived = 0;
                                ReceiveHandler receiver = new ReceiveHandler();
                                do
                                {
                                    channel.Receive(receiver);
                                    receiver.WaitDone();

                                    if (receiver.Success)	// true until ClientSessionChannel is or gets closed
                                    {
                                        byte[] data = receiver.GetData();
                                        if (data != null)
                                        {
                                            totalBytesReceived += data.Length;
                                            if (++receiveCallCnt % 500 == 0)
                                                Console.WriteLine("Thread {0} received {1} bytes so far", m_threadNr, totalBytesReceived);
                                            file.Write(System.Text.Encoding.ASCII.GetString(data).ToCharArray());
                                        }
                                    }
                                } while (receiver.Success && !receiver.Eof());

                                Console.WriteLine("Thread {0} received {1} bytes in total", m_threadNr, totalBytesReceived);
                            }

                            Console.WriteLine("Thread {0} disconnecting for {1}th time", m_threadNr, loop);

                            ph = new ProgressHandler();
                            client.Disconnect(ph);
                            ph.WaitDone();
                        }

                        if (m_coord.m_stopEvent.WaitOne(0))
                            break;
                    }                    
                }
            }
            catch (System.Exception e)
            {
                m_coord.Stop("ExecThread " + m_threadNr + " exited with exception:\r\n" + e.ToString());
            }
        }
    }

    class FlowSshNet_ExecStress
    {
        static void OnUncaughtExceptionInEvent(object sender, bool fatal, System.Exception e)
        {
            if (fatal) Console.WriteLine("Error [fatal] in event: " + e.ToString());
            else       Console.WriteLine("Error in event: " + e.ToString());
            System.Environment.Exit(3);
        }    

        static int Main(string[] args)
        {
            // Register delegate for uncaught exceptions in user-defined events.
            // Example: If there is an uncaught exception in MyClient.OnMyHostKey, 
            //          then this is reported in OnUncaughtExceptionInEvent.
            SshNet.OnExceptionInEvent += new ExceptionInEventEventHandler(OnUncaughtExceptionInEvent);

            if (args.Length < 7)
            {
                Console.WriteLine("Usage: FlowSshNet_ExecStress.exe <host> <port> <hostKeySha256> <user> <pass> <thread1cmdRepeatCnt> <thread1cmd> [; <thread2cmdRepeatCnt> <thread2cmd> ... ]");
                Console.WriteLine("A thread command can be any command you want to execute on the remote server, for instance \"dir\"");
                Console.WriteLine("The returned data of a thread command is stored into a file \"DataRcvdThread_1.txt\", \"DataRcvdThread_2.txt\" ... ");//!text
                Console.WriteLine("Thread commands should not conflict, or have inter-dependencies.");
                Console.WriteLine("Each command will be executed repeatedly and simultaneously on a separate thread.");
                return 2;
            }

            string host = args[0];
            uint   port = uint.Parse(args[1]);
            string hkfp = args[2];
            string user = args[3];
            string pass = args[4];

            List<ExecCommand> commands = new List<ExecCommand>();
            for (int i=5; ; )
            {
                uint cmdRepeatCnt = uint.Parse(args[i]);

                if (++i == args.Length) { Console.WriteLine("Missing command"); return 2; }
                string ecmd = args[i];
                
                ExecCommand cmd = new ExecCommand();
                cmd.m_cmdRepeatCnt = cmdRepeatCnt;
                cmd.m_cmd = ecmd;
                commands.Add(cmd);

                if (++i == args.Length) break;
                if (args[i] != ";")     { Console.WriteLine("Expecting ; between commands"); return 2; }
                if (++i == args.Length) { Console.WriteLine("Expecting command after ; separator"); return 2; }
            }

            Console.WriteLine("Starting threads. Press Esc to exit");

            Coordination coord    = new Coordination();
            List<Thread> threads  = new List<Thread>();
            uint         threadNr = 1;

            foreach (ExecCommand cmd in commands)
            {
                ExecThread xt = new ExecThread();
                xt.m_coord    = coord;
                xt.m_threadNr = threadNr++;
                xt.m_host     = host;
                xt.m_port     = port;
                xt.m_hkfp     = hkfp;
                xt.m_user     = user;
                xt.m_pass     = pass;
                xt.m_cmd      = cmd;

                Thread t = new Thread(xt.Run);
                t.Start();
                threads.Add(t);
            }

            // Wait Esc key or child thread exit
            while (true)
            {
                while (Console.KeyAvailable)
                    if (Console.ReadKey(true).KeyChar == 27)
                    {
                        coord.Stop("Esc key pressed");
                        break;
                    }

                if (IsAnyThreadAlive(threads) == false)
                {
                    coord.Stop("All threads done");
                    break;
                }

                if (coord.m_stopEvent.WaitOne(200))
                    break;
            }            

            Console.WriteLine("Stopping...");
            foreach (Thread t in threads)
                t.Join();

            if (coord.m_stopReason != null)
                Console.WriteLine("Stopped: " + coord.m_stopReason);

            return 0;
        }

        static bool IsAnyThreadAlive(List<Thread> threads)
        {
            foreach (Thread t in threads)
                if (t.IsAlive)
                    return true;
            return false;
        }   
    }
}
