using Newtonsoft.Json;
using System.Data;
using System.Net.WebSockets;

using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;


Console.Write("Enter Jibo IP: ");
var jiboIp = (Console.ReadLine() ?? "").Trim();


if (string.IsNullOrWhiteSpace(jiboIp))
{
    Console.WriteLine("No IP entered.");
    return;
}

var wsAxisState = new Uri($"ws://{jiboIp}:8282/axis_state");
var wsAxisCommand = new Uri($"ws://{jiboIp}:8282/axis_command");

using var cts = new CancellationTokenSource();

Console.WriteLine($"Connecting to Jibo at {jiboIp}...");
Console.WriteLine("Press Ctrl+C to quit.");

Console.CancelKeyPress += (_, e) =>
{
    e.Cancel = true;
    cts.Cancel();
};


while (!cts.IsCancellationRequested)
{
    Console.Write("Enter Jibo new Pose Body <Neck,Torso,Pelvis>" +
        "\r\nDefault is center for all 3 axis 0,0,0" +
        "\r\nUse Degrees separated by a comma 0-359, 0-359, 0-359: ");
    var positionString = (Console.ReadLine() ?? "0,0,0").Trim();
    if (positionString.Length < 1) { positionString = "0,0,0"; }
    float[] position = [.. positionString.Split(',').Select(s => float.TryParse(s.Trim(), out var val) ? val : 0)];

    // According to the web service example,
    // the position values should be in the range of 0-100, where 100 corresponds to 360 degrees.
    if (!string.IsNullOrWhiteSpace(positionString))
    {
        int index = 0;
        Console.WriteLine("Convert to 0-100 scale");
        foreach (var degrees in position)
        {
            position[index++] = (float)((degrees / 3.6) % 100);
        }
    }

    JiboBodyControl js = new();

    try
    {
        js.PositionRobot(position[0], position[1], position[2]);
        js.Start(jiboIp);
    }
    catch (Exception ex)
    {
        Console.WriteLine($"Error: {ex.Message}");
    }


    Console.ReadLine();
    js?.StopAsync().Wait();
}

return;

public struct RobotBodyPart
{
    [JsonProperty(PropertyName = "ts")]
    public List<float>? Ts { get; set; }
    [JsonProperty(PropertyName = "pos")]
    public float Pos { get; set; }
    [JsonProperty(PropertyName = "inc_pos")]
    public float IncPos { get; set; }
    [JsonProperty(PropertyName = "vel")]
    public float Vel { get; set; }
    [JsonProperty(PropertyName = "cur")]
    public float Cur { get; set; }
    [JsonProperty(PropertyName = "pwm")]
    public float Pwm { get; set; }
    [JsonProperty(PropertyName = "status")]
    public int Status { get; set; }
    [JsonProperty(PropertyName = "vel_limit")]
    public float VelLimit { get; set; }
    [JsonProperty(PropertyName = "acc_limit")]
    public float AccLimit { get; set; }
    [JsonProperty(PropertyName = "cur_limit")]
    public float CurLimit { get; set; }
    [JsonProperty(PropertyName = "mode")]
    public int Mode { get; set; }
    [JsonProperty(PropertyName = "ref")]
    public float Ref { get; set; }
    [JsonProperty(PropertyName = "ref_pos")]
    public float RefPos { get; set; }
    [JsonProperty(PropertyName = "ref_acc")]
    public float RefAcc { get; set; }
    [JsonProperty(PropertyName = "ticks")]
    public int Ticks { get; set; }
    [JsonProperty(PropertyName = "integrator")]
    public float Integrator { get; set; }
    [JsonProperty(PropertyName = "fault_status")]
    public int FaultStatus { get; set; }
    [JsonProperty(PropertyName = "index_count")]
    public int IndexCount { get; set; }
    [JsonProperty(PropertyName = "limiter_mode")]
    public int LimiterMode { get; set; }
}

public sealed class BodyState
{
    [JsonProperty(PropertyName = "ts")]
    public List<float>? Ts { get; set; }

    [JsonProperty(PropertyName = "pelvis")]
    public RobotBodyPart Pelvis { get; set; }

    [JsonProperty(PropertyName = "neck")]
    public RobotBodyPart Neck { get; set; }

    [JsonProperty(PropertyName = "torso")]
    public RobotBodyPart Torso { get; set; }

    [JsonProperty(PropertyName = "lockout")]
    public float Lockout { get; set; }

    [JsonProperty(PropertyName = "limiter_torque")]
    public float LimiterTorque { get; set; }

    [JsonProperty(PropertyName = "limiter_lookahead_torque")]
    public float LimiterLookaheadTorque { get; set; }

}

public static class Interval
{
    public static System.Timers.Timer Set(Action action, int interval)
    {
        var timer = new System.Timers.Timer(interval);
        timer.Elapsed += (s, e) =>
        {
            try
            {
                timer.Enabled = false;
                action();
                timer.Enabled = true;
            }
            catch { return; }
        };
        timer.Enabled = true;
        return timer;
    }

    public static void Stop(System.Timers.Timer timer)
    {
        if (timer == null) return;
        timer.Stop();
        timer.Dispose();
    }
}

public sealed class JiboBodyControl
{
    string Host = string.Empty;
    enum Axis { Neck = 0, Torso = 1, Pelvis = 2 }

    const int BREAK = 2;
    const int POS_VEL = 7;
    const int VEL = 4;

    public BodyState? JiboBodyState { get; private set; } = null;

    bool PelvisPositioned = false;
    bool TorsoPositioned = false;
    bool NeckPositioned = false;
    bool Finished = true;

    private float[] PoseArray = new float[3];

    System.Timers.Timer? IndexCheckInterval;
    System.Timers.Timer? SpinInterval;

    Uri? wsAxisState;
    Uri? wsAxisCommand;

    readonly CancellationTokenSource _cts = new();
    ClientWebSocket? _ClientWebSocket = new();
    ClientWebSocket? _ClientControlWebSocket = new();

    public void PositionRobot(float neckPos, float torsoPos, float pelvisPos)
    {
        //pose array, in degrees
        var poseArray_deg = new float[3];
        poseArray_deg[0] = neckPos;
        poseArray_deg[1] = torsoPos;
        poseArray_deg[2] = pelvisPos;

        //pose array, converted to radian
        PoseArray = DegToRad(poseArray_deg);
    }

    private void UpdatePositioning(string host)
    {
        JiboBodyState ??= SetupStateSocketAsync(host).Result;
        UpdateStatus();
    }


    public void Start(string jiboIp)
    {
        PelvisPositioned = false;
        TorsoPositioned = false;
        NeckPositioned = false;
        Finished = false;
        JiboBodyState = null;

        int tryCount = 5;
        Host = jiboIp + ":8282";

        do
        {
            int delay = 1000;
            if (JiboBodyState == null)
            {
                Console.WriteLine("Waiting for initial body state...");
                Task.Delay(delay, _cts.Token).Wait();
                UpdatePositioning(Host);
            }
        }
        while (JiboBodyState == null && !_cts.Token.IsCancellationRequested && --tryCount > 0);

        if (JiboBodyState != null)
        {
            if (SetupCommandSocket(Host).Result)
            {
                Console.WriteLine("Initial body state received. Ready for control loop...");
                ControlLoop();
            }
            else
            {
                Console.WriteLine("Failed to connect to command socket. Control loop will not start.");
                Finished = true;
            }
        }
        else
        {
            Console.WriteLine("Failed to receive initial body state. Control loop will not start.");
            Finished = true;
        }
    }

    public async Task StopAsync()
    {
        if (IndexCheckInterval != null)
        {
            Interval.Stop(IndexCheckInterval);
            IndexCheckInterval = null;
        }
        if (SpinInterval != null)
        {
            Interval.Stop(SpinInterval);
            SpinInterval = null;
        }

        var wsCloseTask = Task.Run(async () =>
        {
            _cts.Cancel();
            if (_ClientWebSocket != null && _ClientWebSocket.State == WebSocketState.Open)
            {
                await _ClientWebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, null, default);
            }
            wsAxisState = null;
            _ClientWebSocket = null;
            Finished = true;
            JiboBodyState = null;
        });

        await Task.WhenAny(wsCloseTask, Task.Delay(TimeSpan.FromSeconds(30), _cts.Token));
        Console.WriteLine("WebCommunications have been shut down\r\nPress enter to set new position.");
    }

    private async Task<bool> SetupCommandSocket(string host)
    {
        if (wsAxisCommand == null)
        {
            wsAxisCommand = new Uri($"ws://{host}/axis_command");
        }
        if (_ClientControlWebSocket == null || _ClientControlWebSocket.State != WebSocketState.Open)
        {
            _ClientControlWebSocket = null;
            _ClientControlWebSocket = new ClientWebSocket();
            _ClientControlWebSocket.Options.KeepAliveInterval = TimeSpan.FromSeconds(10);
        }
        try
        {
            if (_ClientControlWebSocket.State != WebSocketState.Open)
            {
                await _ClientControlWebSocket.ConnectAsync(wsAxisCommand, _cts.Token);
                Console.WriteLine("Command WebSocket connected.");
            }
            else
            {
                Console.WriteLine("Could not Open Command WebSocket.");
                return false;
            }

            return _ClientControlWebSocket != null && _ClientControlWebSocket.State == WebSocketState.Open;
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Couldn\'t connect to body service\'s command websocket: {ex.Message}");
        }

        return false;
    }

    private BodyState? GetBodyState(string json)
    {
        BodyState? bodyState = null;
        var node = JsonNode.Parse(json);
        if (node != null && node?["ts"]?[0]?.GetValue<float>() != 0)
        {
            using (var stringReader = new StringReader(json))
            {
                using var jsonReader = new JsonTextReader(stringReader);
                var serializer = new Newtonsoft.Json.JsonSerializer();
                try
                {
                    bodyState = serializer.Deserialize<BodyState>(jsonReader);
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"Error deserializing JSON: {ex.Message}");
                    bodyState = null;
                }
            }
        }

        return bodyState;
    }

    private async Task<BodyState?> SetupStateSocketAsync(string host)
    {
        if (wsAxisState == null)
        {
            wsAxisState = new Uri($"ws://{host}/axis_state");
        }
        if (_ClientWebSocket == null || _ClientWebSocket.State != WebSocketState.Open)
        {
            _ClientWebSocket = null;
            _ClientWebSocket = new ClientWebSocket();
            _ClientWebSocket.Options.KeepAliveInterval = TimeSpan.FromSeconds(10);
        }

        try
        {
            if (_ClientWebSocket.State != WebSocketState.Open)
            {
                await _ClientWebSocket.ConnectAsync(wsAxisState, _cts.Token);
                Console.WriteLine("WebSocket connected.");
            }

            if (_ClientWebSocket != null && _ClientWebSocket.State == WebSocketState.Open)
            {
                WebSocketReceiveResult result;
                using var ms = new MemoryStream();
                byte[] buffer = new byte[8192];

                var receiveTask = Task.Run(async () =>
                 {
                     int bytesReceived = 0;
                     do
                     {
                         result = await _ClientWebSocket.ReceiveAsync(new ArraySegment<byte>(buffer), _cts.Token);

                         if (result.MessageType == WebSocketMessageType.Close)
                         {
                             // Acknowledge the close to complete the handshake
                             var status = result.CloseStatus ?? WebSocketCloseStatus.NormalClosure;
                             var desc = result.CloseStatusDescription;
                             try
                             {
                                 await _ClientWebSocket.CloseAsync(status, desc, CancellationToken.None);
                             }
                             catch
                             {
                                 // ignore errors during close acknowledgement
                             }
                             wsAxisState = null;
                             break;
                         }

                         ms.Write(buffer, bytesReceived, result.Count);
                         bytesReceived += result.Count;
                     }
                     while (!result.EndOfMessage && wsAxisState != null);

                     if (wsAxisState == null)
                     {
                         Console.WriteLine("WebSocket connection was closed.");
                         JiboBodyState = null;
                         return;
                     }

                     if (result.EndOfMessage)
                     {
                         string json = Encoding.UTF8.GetString(ms.ToArray());
                         JiboBodyState = GetBodyState(json);
                     }
                 });
                
                await receiveTask;

                return JiboBodyState;
            }
            else
            {
                Console.WriteLine("Could not Open WebSocket.");
                JiboBodyState = null;
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Couldn\'t connect to body service\'s state websocket: {ex.Message}");
            // Acknowledge the close to complete the handshake
            var status = WebSocketCloseStatus.NormalClosure;
            var desc = "";
            try
            {
                await _ClientWebSocket.CloseAsync(status, desc, CancellationToken.None);
            }
            catch
            {
                // ignore errors during close acknowledgement
            }
            wsAxisState = null;
            JiboBodyState = null;
        }

        return null;
    }

    private void ControlLoop()
    {
        bool IsPositioning = false;

        if (!_cts.Token.IsCancellationRequested)
        {
            if (JiboBodyState != null)
            {
                if (Finished)
                {
                    Console.WriteLine("Jibo was already positioned.");
                    StopAsync().Wait();
                }
                else
                {
                    if (!IsPositioning)
                    {
                        IsPositioning = true;
                        Console.WriteLine("Positioning robot...");
                        StartSpinningJibo();
                        IndexCheckInterval = Interval.Set(() =>
                        {
                            if (Finished)
                            {
                                StopAsync().Wait();
                            }
                            else
                            {
                                JiboBodyState = null;
                                UpdatePositioning(Host);
                            }

                        }, 30);
                    }
                }
            }
        }
    }

    private async Task<bool> SendCommandAsync(ArraySegment<byte> buffer, string commandJson)
    {
        if (_ClientControlWebSocket != null && _ClientControlWebSocket.State == WebSocketState.Open)
        {
            await _ClientControlWebSocket.SendAsync(buffer, WebSocketMessageType.Text, true, _cts.Token);
            Console.WriteLine("Command sent: " + commandJson);
            return true;
        }

        return false;
    }

    private void StartSpinningJibo()
    {
        if (_ClientControlWebSocket != null && _ClientControlWebSocket.State == WebSocketState.Open)
        {
            SpinInterval = Interval.Set(() => {

                string? commandJson = BuildCommands();
                if (commandJson != null)
                {
                    var bytes = Encoding.UTF8.GetBytes(commandJson);
                    var buffer = new ArraySegment<byte>(bytes);
                    try
                    {
                        SendCommandAsync(buffer,commandJson).Wait();
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine($"Error sending command: {ex.Message}");
                        if (_ClientControlWebSocket?.State == WebSocketState.Closed)
                        {
                            // Acknowledge the close to complete the handshake
                            var status = WebSocketCloseStatus.NormalClosure;
                            var desc = "";
                            try
                            {
                                _ClientControlWebSocket.CloseAsync(status, desc, CancellationToken.None);
                            }
                            catch
                            {
                                // ignore errors during close acknowledgement
                            }
                            StopAsync().Wait();
                            return;
                        }
                    }
                }
            }, 2000);
        }
        else
        {
            Console.WriteLine("Cannot send command, control WebSocket is not open.");
        }
    }

    private float[] DegToRad(float[] poseArray)
    {
        for (int i = 0; i < poseArray.Length; ++i)
        {
            poseArray[i] = float.DegreesToRadians(poseArray[i]);
        }
        return poseArray;
    }

    private float GetAxisRange(Axis axis)
    {
        if (JiboBodyState != null)
        {
            switch (axis)
            {
                case Axis.Pelvis:
                    {
                        var pelvis = JiboBodyState.Pelvis;
                        var mode = pelvis.Mode;
                        float lim = 1;
                        if (mode == 4)
                        {
                            lim = pelvis.VelLimit;
                        }
                        else if (mode == 6)
                        {
                            lim = pelvis.CurLimit;
                        }

                        return lim;
                    }
                case Axis.Torso:
                    {
                        var torso = JiboBodyState.Torso;
                        var mode = torso.Mode;
                        float lim = 1;
                        if (mode == 4)
                        {
                            lim = torso.VelLimit;
                        }
                        else if (mode == 6)
                        {
                            lim = torso.CurLimit;
                        }

                        return lim;
                    }
                case Axis.Neck:
                    {
                        var neck = JiboBodyState.Neck;
                        var mode = neck.Mode;
                        float lim = 1;
                        if (mode == 4)
                        {
                            lim = neck.VelLimit;
                        }
                        else if (mode == 6)
                        {
                            lim = neck.CurLimit;
                        }

                        return lim;
                    }
                default:
                    return 0.0f;
            }
        }

        return 0.0f;
    }

    public void UpdateStatus()
    {
        if (JiboBodyState != null)
        {
            var pelvis = JiboBodyState.Pelvis;
            var torso = JiboBodyState.Torso;
            var neck = JiboBodyState.Neck;

            var tolerance = 0.1;
            PelvisPositioned = (Math.Abs(Math.Abs(pelvis.Pos) -
                Math.Abs(PoseArray[2])) < tolerance);
            TorsoPositioned = (Math.Abs(Math.Abs(torso.Pos) -
                Math.Abs(PoseArray[1])) < tolerance);
            NeckPositioned = (Math.Abs(Math.Abs(neck.Pos) -
                Math.Abs(PoseArray[0])) < tolerance);

            Task.Delay(500, _cts.Token).Wait();

            Finished = PelvisPositioned && TorsoPositioned && NeckPositioned;
        }
        else
        {
            Finished = true;
        }
    }

    private string? BuildCommands()
    {
        if (JiboBodyState == null)
        {
            return null;
        }

        // Grab current body state for command building
        BodyState bodyState = JiboBodyState;

        Console.WriteLine("Starting positioning sequence...");
        string vpelvis;
        string vtorso;
        string vneck;

        float val = GetAxisRange(Axis.Pelvis);
        int mode = VEL;

        if (!PelvisPositioned)
        {
            var currentMode = bodyState.Pelvis.Mode;
            if (currentMode <= BREAK)
            {
                mode = POS_VEL;
            }
        }
        else
        {
           mode = BREAK;
        }
        vpelvis = Cmd(mode, val, PoseArray[2], bodyState.Pelvis.Pos, bodyState.Pelvis);

        val = GetAxisRange(Axis.Torso);
        mode = VEL;

        if (!TorsoPositioned)
        {
            var currentMode = bodyState.Torso.Mode;
            if (currentMode <= BREAK)
            {
                mode = POS_VEL;
            }
        }
        else
        {
           mode = BREAK;
        }
        vtorso = Cmd(mode, val, PoseArray[1], bodyState.Torso.Pos, bodyState.Torso);

        val = GetAxisRange(Axis.Neck);
        mode = VEL;

        if (!NeckPositioned)
        {
            var currentMode = bodyState.Neck.Mode;
            if (currentMode <= BREAK)
            {
                mode = POS_VEL;
            }
        }
        else
        {
           mode = BREAK;
        }
        vneck = Cmd(mode, val, PoseArray[0], bodyState.Neck.Pos, bodyState.Neck);

        StringBuilder cmd = new StringBuilder();
        cmd.Append('{');
        cmd.AppendFormat("\"ts\": [{0}, {1}],", bodyState.Ts?[0], bodyState.Ts?[1]);
        cmd.AppendFormat("\"pelvis\": {0},", vpelvis);
        cmd.AppendFormat("\"torso\": {0},", vtorso);
        cmd.AppendFormat("\"neck\": {0}", vneck);
        cmd.Append('}');

        return cmd.ToString();
    }

     private string Cmd(int mode, float vel, float desPos, float actPos, RobotBodyPart state)
    {
        state.VelLimit = 10;

        StringBuilder cmd = new StringBuilder();
        cmd.Append('{');
        cmd.AppendFormat("\"mode\": {0},", mode);
        cmd.AppendFormat("\"value\": [{0}, {1}],", vel, desPos);
        cmd.AppendFormat("\"vel_limit\": {0},", state.VelLimit);
        cmd.AppendFormat("\"acc_limit\": {0},", state.AccLimit);
        cmd.AppendFormat("\"cur_limit\": {0}", state.CurLimit);
        cmd.Append('}');

        return cmd.ToString();
    }

}