mod codec;
mod pkid;
mod publish_store;

use std::{
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::{anyhow, bail, Result};
use fluent_uri::Uri;
use futures_util::{sink::SinkExt, stream::StreamExt, FutureExt};
use mqttbytes::{v5::*, QoS};
use parking_lot::RwLock;
use tokio::{net::TcpStream, time};
use tokio_util::codec::Framed;
use tracing::{debug, error, warn};

use self::{codec::MqttV5Codec, pkid::PacketId};

use super::{config::MqttClientConfig, Message};

pub(super) use self::publish_store::PublishStore;
use crate::client::config::OnMessageOwnedCallback;

const MAX_MQTT_PACKET_SIZE: usize = 10 * 1024;
const KEEP_ALIVE_INTERVAL: Duration = Duration::from_millis(5_000);
const CONNECTION_TIMEOUT: Duration = Duration::from_millis(5_000);
const RECONNECT_INTERVAL: Duration = Duration::from_millis(1_000);
const STORE_INTERVAL: Duration = Duration::from_millis(1_000);

pub struct MqttWorker {
    mqtt_broker: Uri<String>,
    config: Arc<RwLock<MqttClientConfig>>,
    last_broker_response: Instant,
    publish_store: Arc<PublishStore>,
    on_message_owned: Option<Box<OnMessageOwnedCallback>>,
}

impl MqttWorker {
    pub(super) fn new(
        mqtt_broker: Uri<String>,
        config: Arc<RwLock<MqttClientConfig>>,
        on_message_owned: Option<Box<OnMessageOwnedCallback>>,
    ) -> (Self, Arc<PublishStore>) {
        let publish_store = Arc::new(PublishStore::new());

        let worker = Self {
            mqtt_broker,
            config,
            last_broker_response: Instant::now(),
            publish_store: publish_store.clone(),
            on_message_owned,
        };

        (worker, publish_store)
    }

    /// Blocking MQTT worker main loop, should be spawned in separate tokio::task
    pub async fn run(mut self) {
        let mut interval = time::interval(RECONNECT_INTERVAL);

        loop {
            interval.tick().await;

            // czyścimy ramki QoS = 0 i resetujemy waiting_for_ack
            self.publish_store.on_connect_cleanup();
            let authority = self
                .mqtt_broker
                .authority()
                .expect("missing broker authority");
            let host = authority.host();
            let port = authority
                .port()
                .filter(|p| !p.is_empty())
                .and_then(|p| p.as_str().parse::<u16>().ok())
                .unwrap_or(1883);
            let broker_addr = (host, port);
            debug!("connecting to broker {:?}", broker_addr);
            let socket = match TcpStream::connect(broker_addr).await {
                Ok(s) => s,
                Err(err) => {
                    error!("TCP connection to MQTT broker failed: {}", err);
                    continue;
                }
            };

            if let Err(err) = self.connection_handle(socket).await {
                error!("MQTT connection closed: {}", err);
                continue;
            }
        }
    }

    async fn connection_handle(&mut self, socket: TcpStream) -> Result<()> {
        let mqtt_codec = MqttV5Codec::new(MAX_MQTT_PACKET_SIZE);
        let mut socket = Framed::new(socket, mqtt_codec);

        let _connack = self.mqtt_connect(&mut socket).await?;
        self.on_connected();

        let _suback = self.mqtt_subscribe(&mut socket).await?;

        let mut ping_interval = time::interval(KEEP_ALIVE_INTERVAL);
        let mut store_tick_interval = time::interval(STORE_INTERVAL);

        let mut packet_id = PacketId::new();

        loop {
            tokio::select! {
                // Periodically send Ping packet
                _ = ping_interval.tick().fuse() => {
                    if self.last_broker_response.elapsed() < (KEEP_ALIVE_INTERVAL * 2)
                    {
                        debug!("Sending PingReq");

                        socket.send(&PingReq).await?;
                    } else {
                        bail!("MQTT broker has not responded in {} seconds", self.last_broker_response.elapsed().as_secs())
                    }
                }

                // Handle incoming Packets
                packet = socket.next() => {
                    let received = Instant::now();
                    debug!("MQTT incoming packet: {:?}", packet);
                    let packet = match packet.transpose()? {
                        Some(p) => p,
                        None => {
                            bail!("MQTT broker has finished")
                        }
                    };

                    match packet {
                        Packet::Connect(_connect) => {}
                        Packet::ConnAck(_connack) => {}

                        Packet::Publish(publish) => {
                            let pkid = publish.pkid;
                            let qos = publish.qos;
                            let message = publish.into();

                            self.on_message(&message);
                            if let Some(callback) = &mut self.on_message_owned {
                                callback(message, received);
                            }
                            self.last_broker_response = Instant::now();

                            match qos {
                                QoS::AtMostOnce => (),
                                QoS::AtLeastOnce => {
                                    let puback = PubAck::new(pkid);

                                    debug!("Sending PubAck for packet id = {}", pkid);

                                    // TODO: Resending PubAck messages if there was sending error if it is needed (???)
                                    socket.send(&puback).await?;
                                }
                                QoS::ExactlyOnce => {
                                    // TODO: handle incoming QoS 2 messages
                                    warn!("Received message with QoS 2, missing implementation")
                                },
                            }
                        }

                        Packet::PubAck(puback) => {
                            self.last_broker_response = Instant::now();
                            self.publish_store.remove_waiting_for_ack(puback.pkid);

                            if puback.reason as u8 > 0x00 {
                                warn!("Received PubAck reason code other stan success for pkid = {}: {:?}", puback.pkid, puback.reason);
                            }
                        }

                        Packet::PubRec(_pubrec) => {}
                        Packet::PubRel(_pubrel) => {}
                        Packet::PubComp(_pubcomp) => {}
                        Packet::Subscribe(_subscribe) => {}
                        Packet::SubAck(_suback) => {}
                        Packet::Unsubscribe(_unsubscribe) => {}
                        Packet::UnsubAck(_unsuback) => {}

                        Packet::PingReq => socket.send(&PingResp).await?,

                        Packet::PingResp => {
                            self.last_broker_response = Instant::now()
                        },

                        Packet::Disconnect(disconnect) => {
                            warn!("MQTT disconnected by broker: {:?}", disconnect);

                            return Ok(());
                        }
                    }
                }

                _ = store_tick_interval.tick().fuse() => {
                    self.publish_store.remove_expired();
                },

                _ = self.publish_store.notified().fuse() => {
                    if let Some(publish) = self.publish_store.next_publish(packet_id.into()) {
                        if let Some(0) = publish.properties.as_ref().and_then(|x| x.message_expiry_interval) {
                            debug!("Message timeouted {}", publish.topic);
                            continue;
                        };

                        packet_id.increment();

                        debug!("Message to send: {:?}", publish);
                        socket.send(publish.as_ref()).await?;
                    }
                }
            }
        }
    }

    async fn mqtt_connect(
        &mut self,
        socket: &mut Framed<TcpStream, MqttV5Codec>,
    ) -> Result<ConnAck> {
        let (username, password) = self
            .mqtt_broker
            .authority()
            .and_then(|a| a.userinfo())
            .map(|u| {
                let s = u.as_str();
                match s.split_once(':') {
                    Some((user, pass)) => (user, Some(pass)),
                    None => (s, None),
                }
            })
            .unwrap_or(("", None));
        let login = password.map(|p| Login::new(username, p));

        let properties = ConnectProperties {
            session_expiry_interval: None,
            receive_maximum: None,
            max_packet_size: Some(MAX_MQTT_PACKET_SIZE as u32),
            topic_alias_max: Some(0),
            request_response_info: None,
            request_problem_info: None,
            user_properties: Vec::new(),
            authentication_method: None,
            authentication_data: None,
        };

        let mut connect = { Connect::new(&self.config.read().client_id) };
        let clean_session = { self.config.read().clean_session };

        connect.keep_alive = KEEP_ALIVE_INTERVAL.as_secs() as u16;
        connect.login = login;
        connect.clean_session = clean_session;
        connect.properties = Some(properties);

        socket.send(&connect).await?;

        let rx_packet = time::timeout(CONNECTION_TIMEOUT, async {
            let rx_packet = match socket.next().await {
                Some(packet) => packet?,
                None => bail!("MQTT broker disconnected"),
            };

            let rx_packet = match rx_packet {
                Packet::ConnAck(connack) if connack.code == ConnectReturnCode::Success => connack,
                Packet::ConnAck(connack) => {
                    bail!("Broker rejected. Reason = {:?}", connack.code)
                }
                other => bail!("Expecting ConnAck. Received = {:?}", other),
            };

            Ok::<ConnAck, anyhow::Error>(rx_packet)
        })
        .await??;

        debug!("MQTT client successfully connected to broker");

        self.last_broker_response = Instant::now();

        Ok(rx_packet)
    }

    async fn mqtt_subscribe(
        &mut self,
        socket: &mut Framed<TcpStream, MqttV5Codec>,
    ) -> Result<Option<SubAck>> {
        let mut subscribe = {
            let config = self.config.read();

            if config.subscribe_filters.is_empty() {
                return Ok(None);
            }

            Subscribe::new_many(config.subscribe_filters.iter().cloned())
        };

        let pkid = 1;
        subscribe.pkid = pkid;

        socket.send(&subscribe).await?;

        let rx_packet = time::timeout(CONNECTION_TIMEOUT, async {
            let rx_packet = match socket.next().await {
                Some(packet) => packet?,
                None => bail!("MQTT broker disconnected"),
            };

            let rx_packet = match rx_packet {
                Packet::SubAck(suback) if pkid != suback.pkid => {
                    bail!("Packet Id mismatch")
                }
                Packet::SubAck(suback)
                    if !suback
                        .return_codes
                        .iter()
                        .any(|c| *c as u8 >= SubscribeReasonCode::Unspecified as u8) =>
                {
                    suback
                }
                Packet::SubAck(suback) => {
                    bail!(
                        "Broker rejects subscriptions. Reasons = {:?}",
                        suback.return_codes
                    )
                }
                other => bail!("Expecting SubAck. Received = {:?}", other),
            };

            Ok::<SubAck, anyhow::Error>(rx_packet)
        })
        .await??;

        debug!("MQTT client successfully subscribed topics");

        self.last_broker_response = Instant::now();

        Ok(Some(rx_packet))
    }

    fn on_connected(&self) {
        if let Some(callback) = &self.config.write().on_connected_callback {
            callback()
        }
    }

    fn on_message(&self, message: &Message) {
        if let Some(callback) = &mut self.config.write().on_message_callback {
            callback(message)
        }
    }
}

fn mqttbytes_error_to_anyhow(error: mqttbytes::Error) -> anyhow::Error {
    anyhow!("{}", error)
}
