/* Copyright (c) 2008-2011 - Eric P. Mangold
 * Released under the terms of the MIT/X11 license - see LICENSE.txt */

/* TODO:
 * - Lose the TCP connection when fatal errors occur... follow the semantics of Twisted's implementation.
 *   we currently lose the connection in certain cases, but not others. It's all very ambiguous.
 */
using System;
using System.Text;
using System.Collections;
using System.Collections.Generic;
using System.Threading;
using System.Runtime.CompilerServices;
using System.Windows.Forms;
using System.Net.Sockets;
using Timer=System.Threading.Timer;

// Configure logging for this assembly using the AMPSharp.dll.log4net file
[assembly: log4net.Config.XmlConfigurator(ConfigFileExtension="log4net", Watch=true)]

namespace AMP
{
    public interface IAmpProtocol
    {
        Msg CallRemote(Command cmd, Msg msg);
        IAsyncResult BeginCallRemote(Command cmd, Msg msg, AsyncCallback callback, Object cbState);
        IAsyncResult BeginCallRemote(Command cmd, Msg msg, AsyncCallback callback);
        Msg EndCallRemote(IAsyncResult ar);
        void RegisterResponder(Command cmd, Command.Responder responder, Object callbackState);
        void RegisterResponder(Command cmd, Command.Responder responder);
        void RegisterResponder(Command cmd, Command.Responder responder, Object callbackState, SynchronizationContext callbackContext);
        void RegisterResponder(Command cmd, Command.Responder responder, SynchronizationContext callbackContext);
        void RegisterResponder(Command cmd, Command.Responder responder, Object callbackState, Control callbackControl);
        void RegisterResponder(Command cmd, Command.Responder responder, Control callbackControl);
        void DataReceived(byte[] bytes, int len);
        void ConnectionLost(Exception reason);
        void StartReading();
    }

    /// <summary>
    /// AMP protocol-parsing and state-keeping class.
    /// Instantiate it with a System.IO.Stream (such as System.IO.NetworkStream)
    /// that is already be connected to the remote peer.
    /// Be sure to call StartReading() after you've registered any responders.
    /// Provides CallRemote() and BeginCallRemote(), for making remote AMP calls.
    /// Provides RegisterResponder() for registering delegates that will handle
    /// the AMP commands that we support on this side of the connection. Responders
    /// may be either synchrounous or asynchronous.
    /// </summary>
    public class Protocol: IAmpProtocol
    {
		private static readonly log4net.ILog log = log4net.LogManager.GetLogger("AMP");
        private enum states
        {
            KEY_LEN_READ,
            KEY_DATA_READ,
            VAL_LEN_READ,
            VAL_DATA_READ,
            FULL_MESSAGE_FINISHED
        };

        private states state;
        private ArrayList key_len;
        private ArrayList key_data;
        private ArrayList val_len;
        private int val_len_int;
        private byte[] val_data;

        internal Msg_Raw msg_raw;

        private int lastAskKey;
        private Dictionary<int, Ask_Info> askInfoDict;

        private readonly Dictionary<string, responderStruct> registeredCommands = new Dictionary<string, responderStruct>();

        private IAMPTransport transport;
        private System.IO.Stream stream;
        public System.IO.Stream Stream
        {
            get { return stream; }
        }

        public Protocol()
        {
            state = states.KEY_LEN_READ;
            key_len = new ArrayList();
            key_data = new ArrayList();
            val_len = new ArrayList();
            val_data = new byte[0];

            msg_raw = new Msg_Raw();

            askInfoDict = new Dictionary<int, Ask_Info>();

        }

        public Protocol(System.IO.Stream stream)
            : this()
        {
            this.stream = stream;
            transport = new StreamTransport(stream);
            transport.RegisterProtocol(this);
        }

        public void StartReading()
        {
            transport.StartReading();
        }

        // Hold a reference to an Responder, the Command for this responder and the responder state object
        private struct responderStruct
        {
            public Command.Responder responder;
            public Command command;
            public Object state; // may be null
            public SynchronizationContext callbackContext; // may be null
            public Control callbackControl; // may be null
        }

        public class CommandAlreadyRegistered: Exception { };

        public void RegisterResponder(Command cmd, Command.Responder responder)
        {
            _RegisterResponder(cmd, responder, null, null, null);
        }

        public void RegisterResponder(Command cmd, Command.Responder responder, object callbackState, SynchronizationContext callbackContext)
        {
            _RegisterResponder(cmd, responder, callbackState, callbackContext, null);
        }

        public void RegisterResponder(Command cmd, Command.Responder responder, SynchronizationContext callbackContext)
        {
            _RegisterResponder(cmd, responder, null, callbackContext, null);
        }

        public void RegisterResponder(Command cmd, Command.Responder responder, object callbackState, Control callbackControl)
        {
            _RegisterResponder(cmd, responder, callbackState, null, callbackControl);
        }

        public void RegisterResponder(Command cmd, Command.Responder responder, Control callbackControl)
        {
            _RegisterResponder(cmd, responder, null, null, callbackControl);
        }
        public void RegisterResponder(Command cmd, Command.Responder responder, Object callbackState)
        {
            _RegisterResponder(cmd, responder, callbackState, null, null);
        }
        private void _RegisterResponder(Command cmd, Command.Responder responder, Object callbackState,
            SynchronizationContext callbackContext, Control callbackControl)
        {
            if (registeredCommands.ContainsKey(cmd.Name))
            {
                throw new CommandAlreadyRegistered();
            }
            var respRef = new responderStruct {responder = responder,
                command = cmd, state = callbackState, callbackContext = callbackContext,
                callbackControl = callbackControl};
            registeredCommands.Add(cmd.Name, respRef);
        }

        public virtual void ConnectionLost(Exception reason)
        {
            foreach (Ask_Info info in askInfoDict.Values)
            {
                info.ar.error = reason;
                try
                {
                    info.ar.doCallback();
                }
                catch (Exception e)
                {
                    log.Error("Exception in user-defined AMP callback while failing outstanding requests:", e);
                }
            }

            // dont bother to do this if we are gonna be closing the stream associated with
            // this protocol.. 
            askInfoDict.Clear();
        }

        public void DataReceived(byte[] buf)
        {
            DataReceived(buf, buf.Length);
        }

        [MethodImpl(MethodImplOptions.Synchronized)]
        public void DataReceived(byte[] buf, int len)
        {
            int idx = 0;

            while (idx < len)
            {
                int bytes_left = len - idx;
                int bytes_needed;
                switch (state)
                {
                    case states.KEY_LEN_READ:
                        if (key_len.Count == 0 && buf[idx] != 0)
                        {
                            // The 1st byte of the key length must always be \x00
                            throw new Error.ProtocolError();
                        }
                        key_len.Add(buf[idx]);
                        idx++;
                        if (key_len.Count == 2) {
                            // got the 2 bytes for a key. transition state.
                            // if these 2 bytes are both null, then that should signal the
                            // end of a full AMP message.
                            if ((byte)key_len[0] == 0 && (byte)key_len[1] == 0) {
                                if (msg_raw.Count == 0) {
                                    // we haven't received any key/values yet so there shouldn't be a 
                                    // double-NULL here
                                    throw new Error.ProtocolError();
                                }
                                // We got the empty-key terminator on a full, valid message, yay!
                                key_len.Clear(); // forget the 2 null bytes
                                ProcessFullMessage(msg_raw);
                                state = states.FULL_MESSAGE_FINISHED; // transition
                            } else {
                                state = states.KEY_DATA_READ; // transition
                            }
                        } // else state stays the same since we've only read 1 byte of the 2-byte key length
                        break;

                    case states.KEY_DATA_READ:
                        // the key length prefix is restricted to 1 byte (8 bits) on the wire, even though it is
                        // represented by two bytes. Thus, because it is big-endian, the first byte will always be zero
                        // and the 2nd byte is the only one we are interested in.
                        int key_len_int = (byte)key_len[1];
                        
                        bytes_needed = key_len_int - key_data.Count;
                        byte[] tmp;
                        if (bytes_left < bytes_needed) {
                            tmp = new byte[bytes_left];
                            Array.Copy(buf, idx, tmp, 0, bytes_left);
                            key_data.AddRange(tmp);
                            idx += bytes_left;
                            // state stays the same
                        } else if (bytes_left >= bytes_needed) {
                            tmp = new byte[bytes_needed];
                            Array.Copy(buf, idx, tmp, 0, bytes_needed);
                            key_data.AddRange(tmp);
                            idx += bytes_needed;
                            state = states.VAL_LEN_READ; // transition
                        }
                        break;

                    case states.VAL_LEN_READ:
                        val_len.Add(buf[idx]);
                        idx++;
                        if (val_len.Count == 2) {
                            // TODO use a XOR SWAP instead of a separate tmp_bytes array

                            // got the 2 bytes for a value length.
                            var tmp_bytes = new byte[2];
                            if (BitConverter.IsLittleEndian)
                            {
                                // reverse bytes before casting because this is a little-endian computer, but the bytes
                                // we read are in network byte order (big-endian)
                                tmp_bytes[0] = (byte)val_len[1];
                                tmp_bytes[1] = (byte)val_len[0];
                            }
                            else
                            {
                                tmp_bytes[0] = (byte)val_len[0];
                                tmp_bytes[1] = (byte)val_len[1];
                            }
                            val_len_int = BitConverter.ToUInt16(tmp_bytes, 0);

                            state = states.VAL_DATA_READ; // transition
                        }
                        // else the state stays the same since we've only read 1 byte of the 2-byte value length
                        break;

                    case states.VAL_DATA_READ:
                        bytes_needed = val_len_int - val_data.Length;
                        if (bytes_left < bytes_needed) {
                            var tmp_bytes = new byte[val_data.Length+bytes_left];
                            Array.Copy(val_data, tmp_bytes, val_data.Length);
                            Array.Copy(buf, idx, tmp_bytes, val_data.Length, bytes_left);
                            val_data = tmp_bytes;
                            idx += bytes_left;
                            // state stays the same
                        } else if (bytes_left >= bytes_needed) {
                            var tmp_bytes = new byte[val_data.Length+bytes_needed];
                            Array.Copy(val_data, tmp_bytes, val_data.Length);
                            Array.Copy(buf, idx, tmp_bytes, val_data.Length, bytes_needed);
                            val_data = tmp_bytes;
                            idx += bytes_needed;
                            state = states.KEY_LEN_READ; // transition
                            ProcessCurrentKeyValuePair();
                        }
                        break;

                    case states.FULL_MESSAGE_FINISHED:
                        // clear out the msg_raw map so we can fill it w/ the next one
                        msg_raw = new Msg_Raw();
                        state = states.KEY_LEN_READ;
                        break;
                }
            }
        }

        internal virtual void ProcessFullMessage(Msg_Raw raw_msg)
        {
            if ( raw_msg.ContainsKey("_command") )
            {
                string command = Encoding.UTF8.GetString(raw_msg["_command"]);
                ProcessCommand(command, raw_msg);
            } else if ( raw_msg.ContainsKey("_answer") ) {
                ProcessAnswer(raw_msg);
            } else if ( raw_msg.ContainsKey("_error") ) {
                ProcessError(raw_msg);
            } else {
                log.Error("Got AMP message w/o a _command, _answer, or _error key!");
                return;
            }
        }

        private void ProcessCommand(string cmdName, Msg_Raw raw_msg)
        {
            log.Debug("Processing AMP command: " + cmdName);

            if (registeredCommands.ContainsKey(cmdName))
            {
                // YES, THIS COMMAND HAS A REGISTERED RESPONDER.
                responderStruct respRef = registeredCommands[cmdName];
                Msg requestMsg = respRef.command.ToTypedMessage(raw_msg, MsgType.REQUEST);
                try
                {
                    Object returnVal = null;
                    if (respRef.callbackContext != null)
                    {
                        respRef.callbackContext.Send(delegate
                             {
                                 returnVal = respRef.responder(requestMsg, respRef.state);
                             }, null);
                    }
                    else if (respRef.callbackControl != null)
                    {
                        returnVal = respRef.callbackControl.Invoke(respRef.responder, requestMsg, respRef.state);
                    }
                    else
                    {
                        returnVal = respRef.responder(requestMsg, respRef.state);
                    }

                    if (!raw_msg.ContainsKey("_ask")) return;

                    var responseMsg = returnVal as Msg;
                    if (responseMsg != null)
                    {
                        SendAMPResponse(responseMsg, respRef.command, raw_msg["_ask"]);
                    }
                    else
                    {
                        var deferred = returnVal as DeferredResponse;
                        if (deferred != null)
                        {
                            byte[] ask_key = raw_msg["_ask"];
                            deferred.setCallback(response => ProcessDeferredAMPResponse(response, respRef.command, ask_key));
                        }
                        else
                        {
                            throw new Error.BadResponderReturnValue();
                        }
                    }
                }
                catch (Exception e)
                {
                    // TODO FIXME - SendAMPError does error logging, but we should still
                    // log the error even if no _ask key.
                    if (raw_msg.ContainsKey("_ask"))
                    {
                        SendAMPError(e, respRef.command, raw_msg["_ask"]);
                    }
                }
            }
            else
            {
                // NO, THERE ISN'T ANY REGISTERED RESPONDER FOR THIS COMMAND
                log.Error("No such AMP command: " + cmdName);
                if (raw_msg.ContainsKey("_ask"))
                {
                    var response = new Msg_Raw
                                       {
                                           {"_error", raw_msg["_ask"]},
                                           {"_error_code", Error.Codes.UNHANDLED_ERROR_CODE},
                                           {
                                               "_error_description",
                                               Encoding.UTF8.GetBytes(System.String.Format("Unhandled Command: {0}", cmdName))
                                               }
                                       };
                    SendMessage(response);
                }
                return;
            }
        }

        private void SendAMPError(Exception e, Command cmd, byte[] ask_key)
        {
            var response = new Msg_Raw {{"_error", ask_key}};
            byte[] error_code;
            byte[] descr;
            if (cmd.errors.ContainsKey(e.GetType()))
            {
                error_code = cmd.errors[e.GetType()];
                descr = Encoding.UTF8.GetBytes(e.ToString());
            }
            else
            {
                log.Error("Exception occured in command handler that was not defined on Command object:");
                log.Error(e.ToString());
                error_code = Error.Codes.UNKNOWN_ERROR_CODE;
                descr = Encoding.UTF8.GetBytes("Unknown Error");
            }
            response.Add("_error_code", error_code);
            response.Add("_error_description", descr);
            SendMessage(response);
        }

        private void SendAMPResponse(Msg msg, Command cmd, byte[] answer_key)
        {
            // XXX TODO - should log a helpful error message if this raises InvalidCastException
            // because it's very easy for the user to have put a wrong type in to the Msg instance.
            Msg_Raw responseMsg = cmd.ToRawMessage(msg, MsgType.RESPONSE);
            responseMsg["_answer"] = answer_key;
            SendMessage(responseMsg);
        }

        internal virtual void SendMessage(Msg_Raw raw_msg)
        {
            transport.Write(BuildAMPWireCommand(raw_msg));
        }

        private void ProcessDeferredAMPResponse(Object response, Command cmd, byte[] answer_key)
        {
            try
            {
                try
                {
                    var msg = (Msg)response;
                    SendAMPResponse(msg, cmd, answer_key);
                }
                catch (InvalidCastException)
                {
                    // assume it's an exception
                    var err = (Exception) response;
                    SendAMPError(err, cmd, answer_key);
                }
            }
            catch (Exception e)
            {
                SendAMPError(e, cmd, answer_key);
            }
        }

        private void ProcessAnswer(Msg_Raw raw_msg)
        {
            int answerKey = Int32.Parse(Encoding.UTF8.GetString(raw_msg["_answer"]));
            log.Debug("Processing AMP answer for request " + answerKey);
            if ( ! askInfoDict.ContainsKey(answerKey) ) {
                log.ErrorFormat("Got AMP answer message w/o an unknown key {0}! (Perhaps we already timed out?)", answerKey);
                return;
            }
            Ask_Info askInfo = askInfoDict[answerKey];
            askInfo.ar.result = askInfo.cmd.ToTypedMessage(raw_msg, MsgType.RESPONSE);
            try
            {
                //askInfo.ar._callback(askInfo.ar);
                askInfo.ar.doCallback();

            } catch (Exception e) {
                // XXX TODO should this catch exceptions? We aren't really writing a *framework*, a la Twisted, here. So since we're just
                // writing a *library*, maybe it should be up to the client code to deal with exceptions.
                // NOTE: We can't catch errors anyway in the case that doCallback posts the callback on the SyncronizationContext
                log.Error("Exception raised in user-defined AMP result callback:", e);
            } finally {
                askInfoDict.Remove(answerKey);
            }
        }

        private void ProcessError(Msg_Raw raw_msg)
        {
            // According to the AMP spec an error packet should contain all 3 of these keys
            int msgKey = Int32.Parse(Encoding.UTF8.GetString(raw_msg["_error"]));

            byte[] errorCode = raw_msg["_error_code"];
            string description = Encoding.UTF8.GetString(raw_msg["_error_description"]);

            log.DebugFormat("Procesing AMP error for request {0} code: {1} descr: {2}", msgKey, BytesUtil.ToHexStr(errorCode), description);

            if ( ! askInfoDict.ContainsKey(msgKey) ) {
                log.ErrorFormat("Got AMP error message w/ an unknown key {0}! (Perhaps we already timed out?)", msgKey);
                return;
            }

            Ask_Info askInfo = askInfoDict[msgKey];

            System.Type excType = typeof(Error.UnknownRemoteError); // default if we don't find the _error_code below

            foreach (var kvp in askInfo.cmd.errors)
            {
                if (BytesUtil.AreEqual(kvp.Value, errorCode))
                {
                    excType = kvp.Key;
                }
            }

            Exception exc = null;
            try
            {
                // If excType is a subclass of RemoteAmpError then we should find a
                // matching constructor here.
                exc = (Exception)Activator.CreateInstance(excType, errorCode, description);
            }
            catch (MissingMethodException)
            {
                // Otherwise any type of Exception should at least have a void constructor
                exc = (Exception)Activator.CreateInstance(excType);
            }

            askInfo.ar.error = exc;
            try
            {
                askInfo.ar.doCallback();
            } catch (Exception e) {
                // XXX TODO should this catch exceptions? We aren't really writing a *framework*, a la Twisted, here. So since we're just
                // writing a *library*, maybe it should be up to the client code to deal with exceptions.
                log.Error("AMP error returned, and exception raised in user-defined callback:", e);
            } finally {
                askInfoDict.Remove(msgKey);
            }
        }

        private void ProcessCurrentKeyValuePair()
        {
            var key_data_array = (byte[])key_data.ToArray(typeof(byte));
            msg_raw.Add(Encoding.UTF8.GetString(key_data_array), val_data);

            key_len.Clear();
            key_data.Clear();

            val_len.Clear();
            val_data = new byte[0];

            val_len_int = 0;
        }

        // TODO Benchmark this method and experiment with different algorithms
        internal static byte[] BuildAMPWireCommand(Msg_Raw msg)
        {
            var tmp_result = new ArrayList();
            byte[] value;

            if (msg.Count == 0)
            {
                throw new Error.MessageEmptyError();
            }
            if (BitConverter.IsLittleEndian)
            {
                foreach (string key in msg.Keys)
                {
                    value = msg[key];
                    if (key.Length == 0 || key.Length > 255) throw new Error.InvalidKeySizeError();
                    if (value.Length > 65535) throw new Error.InvalidValueSizeError();
                    tmp_result.Add((byte)0x00); // first byte of key length
                    // use the first byte of the 2-byte sequence, and store that as the 2nd byte in the tmp_result
                    // this is because the host computer is little-endien, but we're encoding wire values in network byte order (big-endien)
                    tmp_result.Add(BitConverter.GetBytes((UInt16)key.Length)[0]);
                    tmp_result.AddRange(Encoding.UTF8.GetBytes(key));
                    byte[] tmp = BitConverter.GetBytes((UInt16)value.Length);
                    tmp_result.Add(tmp[1]);
                    tmp_result.Add(tmp[0]);
                    tmp_result.AddRange(value);
                }
            }
            else
            {
                foreach (string key in msg.Keys)
                {
                    value = msg[key];
                    if (key.Length == 0 || key.Length > 255) throw new Error.InvalidKeySizeError();
                    if (value.Length > 65535) throw new Error.InvalidValueSizeError();
                    tmp_result.Add((byte)0x00); // first byte of key length

                    tmp_result.Add(BitConverter.GetBytes((UInt16)key.Length)[1]);
                    tmp_result.AddRange(Encoding.UTF8.GetBytes(key));
                    byte[] tmp = BitConverter.GetBytes((UInt16)value.Length);
                    tmp_result.Add(tmp[0]);
                    tmp_result.Add(tmp[1]);
                    tmp_result.AddRange(value);
                }
            }
            // terminated w/ an empty key
            tmp_result.Add( (byte)0x00 );
            tmp_result.Add( (byte)0x00 );

            //byte[] result = {1};
            //return (byte[])tmp_result;
            //return tmp_result.ToArr
            var result = (byte[])tmp_result.ToArray(typeof(byte));
            return result;
        }

        public Msg CallRemote(Command cmd, Msg msg)
        {
            IAsyncResult ar = BeginCallRemote(cmd, msg, null, null);
            ar.AsyncWaitHandle.WaitOne();
            return EndCallRemote(ar);
        }

        public IAsyncResult BeginCallRemote(Command cmd, Msg msg, AsyncCallback callback)
        {
            return BeginCallRemote(cmd, msg, callback, null);
        }

        [MethodImpl(MethodImplOptions.Synchronized)]
        public IAsyncResult BeginCallRemote(Command cmd, Msg msg, AsyncCallback callback, Object cbState)
        {
            var ar = new SimpleAsyncResult {callback = callback, asyncState = cbState};
            Msg_Raw raw_msg = cmd.ToRawMessage(msg, MsgType.REQUEST);
            raw_msg.Add("_command", Encoding.UTF8.GetBytes(cmd.Name));
            if (cmd.RequiresAnswer)
            {
                int currentKey = ++lastAskKey;
                askInfoDict.Add(currentKey, new Ask_Info(currentKey, cmd, ar));
                raw_msg.Add("_ask", Encoding.UTF8.GetBytes(currentKey.ToString()));
                log.DebugFormat("Calling remote AMP Command: {0} with key {1}", cmd.Name, currentKey);
                SendMessage(raw_msg);
            }
            else
            {
                log.DebugFormat("Calling remote AMP Command: {0}", cmd.Name);
                SendMessage(raw_msg);
                ar.doCallback();
            }
            return ar;
        }

        public Msg EndCallRemote(IAsyncResult ar)
        {
            SimpleAsyncResult result = ThrowIfError(ar);

            // No error, we must have a result from the server available, then
            // If RequiresAnswer was false on the Command that was called, then
            // result should simply be null here.
            return (Msg)result.result;
        }

        private static SimpleAsyncResult ThrowIfError(IAsyncResult ar)
        {
            var result = (SimpleAsyncResult)ar;
            if (result.error != null)
            {
                throw result.error;
            }
            return result;
        }
    }

    public class SimpleAsyncResult : IAsyncResult
    {
        internal object result;
        internal Exception error;
        
        // TODO make priviate/internal (needed by tests, though)
        public AsyncCallback callback;
        internal SynchronizationContext savedContext;

        public SimpleAsyncResult()
        {
            savedContext = SynchronizationContext.Current;
            waitHandle = new ManualResetEvent(false);
        }

        public void ThrowIfError()
        {
            if (error != null)
            {
                throw error;
            }
        }

        // TODO make priviate (needed by tests, though)
        public void doCallback()
        {
            isCompleted = true;
            try
            {
                if (callback != null)
                {
                    if (savedContext != null)
                    {
                        // Post the callback to the context
                        // (Callback will happen asynchronously in the saved context's thread)
                        // XXX TODO shouldn't we pass _asyncState to the delegate here, and also
                        // pass it to the callback in the case where we run the callback synchronously (below) ?
                        // maybe not, because AsyncCallback is only defined to take a single argument: IAsyncResult
                        savedContext.Post(delegate { callback(this); }, null);
                    }
                    else
                    {
                        // Just callback synchronously
                        callback(this);
                    }
                }
            }
            finally
            {
                waitHandle.Set();
            }
        }

        // BEGIN IAsyncResult properties

        // TODO these really should be defined 'internal'
        public bool isCompleted;
        public Object asyncState;
        public bool completedSynchronously;
        public ManualResetEvent waitHandle;

        public Object AsyncState
        {
            get
            {
                return asyncState;
            }
        }

        public WaitHandle AsyncWaitHandle
        {
            get
            {
                return waitHandle;
            }
        }

        public bool IsCompleted {
            get
            {
                return isCompleted;
            }
        }
        public bool CompletedSynchronously {
            get
            {
                return completedSynchronously;
            }
        }

        // END IAsyncResult properties

    }
}
