use anyhow::anyhow;
use api::data_type::DataType;
use api::data_value::DataValue;
use api::request::{
    RequestMessage, RequestMessagePayload, TelemetryDefinitionRequest, TelemetryEntry,
};
use api::response::{ResponseMessage, ResponseMessagePayload, TelemetryDefinitionResponse};
use chrono::{DateTime, TimeDelta, Utc};
use futures_util::future::join_all;
use futures_util::{SinkExt, StreamExt};
use num_traits::FloatConst;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::sync::mpsc::channel;
use tokio::time::{sleep, sleep_until, Instant};
use tokio::{pin, select};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

struct Telemetry {
    is_connected: Arc<AtomicBool>,
    tx: mpsc::Sender<RequestMessage>,
    callback_new: mpsc::Sender<(Uuid, mpsc::Sender<ResponseMessagePayload>)>,
    callback_delete: mpsc::Sender<Uuid>,
    cancel: CancellationToken,
}

struct TelemetryItemHandle<'a> {
    uuid: Uuid,
    tlm: &'a Telemetry,
}

impl Telemetry {
    async fn connect<R>(request: R) -> anyhow::Result<WebSocketStream<MaybeTlsStream<TcpStream>>>
    where
        R: IntoClientRequest + Unpin,
    {
        let (client, _) = match connect_async(request).await {
            Ok(o) => o,
            Err(e) => {
                sleep(Duration::from_secs(1)).await;
                return Err(e.into());
            }
        };
        Ok(client)
    }

    pub async fn new<R>(request: R) -> anyhow::Result<Self>
    where
        R: IntoClientRequest + Unpin + Clone + Send + 'static,
    {
        let cancel = CancellationToken::new();
        let cancel_stored = cancel.clone();
        let (tx_tx, mut tx_rx) = channel(128);
        let (callback_new_tx, mut callback_new_rx) = channel(16);
        let (callback_delete_tx, mut callback_delete_rx) = channel(16);

        let is_connected = Arc::new(AtomicBool::new(false));
        let connected = is_connected.clone();

        tokio::spawn(async move {
            while !cancel.is_cancelled() {
                let client = Self::connect(request.clone());
                pin!(client);
                let mut client = match loop {
                    break select! {
                        c = &mut client => {
                            c
                        },
                        Some(_) = tx_rx.recv() => {continue;},
                        Some(_) = callback_new_rx.recv() => {continue;},
                        Some(_) = callback_delete_rx.recv() => {continue;},
                    };
                } {
                    Ok(c) => c,
                    Err(e) => {
                        println!("Connect Error: {e}");
                        continue;
                    }
                };

                let mut close_connection = true;
                let mut callbacks = HashMap::<Uuid, mpsc::Sender<ResponseMessagePayload>>::new();
                connected.store(true, Ordering::SeqCst);

                loop {
                    select! {
                        biased;
                        _ = cancel.cancelled() => { break; },
                        Some(msg) = callback_new_rx.recv() => {
                            let (uuid, callback) = msg;
                            callbacks.insert(uuid, callback);
                        },
                        Some(msg) = callback_delete_rx.recv() => {
                            callbacks.remove(&msg);
                        },
                        Some(msg) = client.next() => {
                            match msg {
                                Ok(msg) => {
                                    match msg {
                                        Message::Text(msg) => {
                                            let msg: ResponseMessage = match serde_json::from_str(&msg) {
                                                Ok(m) => m,
                                                Err(e) => {
                                                    println!("Failed to deserialize {e}");
                                                    break;
                                                }
                                            };
                                            if let Some(cb) = callbacks.get_mut(&msg.uuid) {
                                                if let Err(e) = cb.send(msg.payload).await {
                                                    println!("Failed to call callback {e}");
                                                    callbacks.remove(&msg.uuid);
                                                }
                                            } else {
                                                unimplemented!("Unexpected Message: {msg:?}");
                                            }
                                        }
                                        Message::Binary(_) => unimplemented!("Binary Unsupported"),
                                        Message::Ping(data) => {
                                            if let Err(e) = client.send(Message::Pong(data)).await {
                                                println!("Failed to send Pong {e}");
                                                break;
                                            }
                                        }
                                        Message::Pong(_) => {
                                            // Intentionally Left Empty
                                        }
                                        Message::Close(_) => {
                                            println!("Websocket Closed");
                                            close_connection = false;
                                            break;
                                        }
                                        Message::Frame(_) => unreachable!("Not Possible"),
                                    }
                                }
                                Err(e) => {
                                    println!("Receive Error {e}");
                                    break;
                                }
                            }
                        },
                        Some(msg) = tx_rx.recv() => {
                            let msg = match serde_json::to_string(&msg) {
                                Ok(m) => m,
                                Err(e) => {
                                    println!("Encode Error {e}");
                                    break;
                                }
                            };
                            if let Err(e) = client.send(Message::Text(msg.into())).await {
                                println!("Send Error {e}");
                                break;
                            }
                        },
                        else => { break; },
                    }
                }
                connected.store(false, Ordering::SeqCst);
                if close_connection {
                    if let Err(e) = client.close(None).await {
                        println!("Close Error {e}");
                    }
                }

                // Callbacks gets dropped here - closing all listeners automatically
            }
        });

        Ok(Self {
            tx: tx_tx,
            callback_new: callback_new_tx,
            callback_delete: callback_delete_tx,
            cancel: cancel_stored,
            is_connected,
        })
    }

    pub async fn wait_connected(&self) {
        while !self.is_connected.load(Ordering::Relaxed) {
            sleep(Duration::from_millis(10)).await;
        }
    }

    /// Send a message and don't expect a response
    pub async fn send_message<P: Into<RequestMessagePayload>>(
        &self,
        payload: P,
    ) -> anyhow::Result<()> {
        self.tx
            .send(RequestMessage {
                uuid: Uuid::new_v4(),
                payload: payload.into(),
            })
            .await?;
        Ok(())
    }

    async fn send_request<
        P: Into<RequestMessagePayload>,
        R: TryFrom<ResponseMessagePayload, Error = impl std::error::Error + Send + Sync + 'static>,
    >(
        &self,
        payload: P,
    ) -> anyhow::Result<R> {
        let uuid = Uuid::new_v4();
        let (tx, mut rx) = channel(1);

        self.wait_connected().await;
        self.callback_new.send((uuid, tx)).await?;
        self.tx
            .send(RequestMessage {
                uuid,
                payload: payload.into(),
            })
            .await?;

        let response = rx.recv().await.ok_or(anyhow!("No Response Received"))?;

        let response = R::try_from(response)?;

        self.callback_delete.send(uuid).await?;

        Ok(response)
    }

    pub async fn register(
        &self,
        name: String,
        data_type: DataType,
    ) -> anyhow::Result<TelemetryItemHandle<'_>> {
        let response: TelemetryDefinitionResponse = self
            .send_request(TelemetryDefinitionRequest { name, data_type })
            .await?;

        Ok(TelemetryItemHandle {
            uuid: response.uuid,
            tlm: self,
        })
    }
}

impl Drop for Telemetry {
    fn drop(&mut self) {
        self.cancel.cancel();
    }
}

impl<'a> TelemetryItemHandle<'a> {
    pub async fn publish(&self, value: DataValue, timestamp: DateTime<Utc>) -> anyhow::Result<()> {
        self.tlm
            .send_message(TelemetryEntry {
                uuid: self.uuid,
                value,
                timestamp,
            })
            .await?;
        Ok(())
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let tlm = Telemetry::new("ws://[::1]:8080/backend").await?;

    {
        let time_offset = tlm
            .register("simple_producer/time_offset".into(), DataType::Float64)
            .await?;

        let publish_offset = tlm
            .register("simple_producer/publish_offset".into(), DataType::Float64)
            .await?;

        let await_offset = tlm
            .register("simple_producer/await_offset".into(), DataType::Float64)
            .await?;

        let sin_tlm_handle = tlm
            .register("simple_producer/sin".into(), DataType::Float32)
            .await?;
        let cos_tlm_handle = tlm
            .register("simple_producer/cos".into(), DataType::Float64)
            .await?;
        let bool_tlm_handle = tlm
            .register("simple_producer/bool".into(), DataType::Boolean)
            .await?;

        let sin2_tlm_handle = tlm
            .register("simple_producer/sin2".into(), DataType::Float32)
            .await?;
        let cos2_tlm_handle = tlm
            .register("simple_producer/cos2".into(), DataType::Float64)
            .await?;

        let sin3_tlm_handle = tlm
            .register("simple_producer/sin3".into(), DataType::Float32)
            .await?;
        let cos3_tlm_handle = tlm
            .register("simple_producer/cos3".into(), DataType::Float64)
            .await?;

        let sin4_tlm_handle = tlm
            .register("simple_producer/sin4".into(), DataType::Float32)
            .await?;
        let cos4_tlm_handle = tlm
            .register("simple_producer/cos4".into(), DataType::Float64)
            .await?;

        let sin5_tlm_handle = tlm
            .register("simple_producer/sin5".into(), DataType::Float32)
            .await?;
        let cos5_tlm_handle = tlm
            .register("simple_producer/cos5".into(), DataType::Float64)
            .await?;

        let sin6_tlm_handle = tlm
            .register("simple_producer/sin6".into(), DataType::Float32)
            .await?;
        let cos6_tlm_handle = tlm
            .register("simple_producer/cos6".into(), DataType::Float64)
            .await?;

        let cancellation_token = CancellationToken::new();
        {
            let cancellation_token = cancellation_token.clone();
            tokio::spawn(async move {
                let _ = tokio::signal::ctrl_c().await;
                cancellation_token.cancel();
                println!("Cancellation Token Cancelled");
            });
        }

        let start_time = Utc::now();
        let start_instant = Instant::now();
        let mut next_time = start_instant;
        let mut index = 0;
        let mut tasks = vec![];
        while !cancellation_token.is_cancelled() {
            next_time += Duration::from_millis(10);
            index += 1;
            sleep_until(next_time).await;
            let publish_time = start_time + TimeDelta::from_std(next_time - start_instant).unwrap();
            let actual_time = Instant::now();
            tasks.push(time_offset.publish(
                DataValue::Float64((actual_time - next_time).as_secs_f64()),
                Utc::now(),
            ));
            tasks.push(sin_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (1000.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (1000.0_f64)).cos()),
                publish_time,
            ));
            tasks.push(
                bool_tlm_handle.publish(DataValue::Boolean(index % 1000 > 500), publish_time),
            );
            tasks.push(sin2_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (500.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos2_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (500.0_f64)).cos()),
                publish_time,
            ));
            tasks.push(sin3_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (333.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos3_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (333.0_f64)).cos()),
                publish_time,
            ));
            tasks.push(sin4_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (250.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos4_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (250.0_f64)).cos()),
                publish_time,
            ));
            tasks.push(sin5_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (200.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos5_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (200.0_f64)).cos()),
                publish_time,
            ));
            tasks.push(sin6_tlm_handle.publish(
                DataValue::Float32((f32::TAU() * (index as f32) / (166.0_f32)).sin()),
                publish_time,
            ));
            tasks.push(cos6_tlm_handle.publish(
                DataValue::Float64((f64::TAU() * (index as f64) / (166.0_f64)).cos()),
                publish_time,
            ));

            tasks.push(publish_offset.publish(
                DataValue::Float64((Instant::now() - actual_time).as_secs_f64()),
                Utc::now(),
            ));

            // Join the tasks so they all run in parallel
            for task in join_all(tasks.drain(..)).await {
                task?;
            }

            tasks.push(await_offset.publish(
                DataValue::Float64((Instant::now() - actual_time).as_secs_f64()),
                Utc::now(),
            ));
        }
        println!("Exiting Loop");
    }
    drop(tlm);

    Ok(())
}