use crate::protocol::codec::frame::FrameType;
use bytes::Bytes;
use dashmap::DashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc::{error::SendError, Sender};

pub struct IdGenerator {
    counter: AtomicU32,
}

impl IdGenerator {
    pub fn new(is_client: bool) -> Self {
        let start = if is_client { 1 } else { 2 };
        Self {
            counter: AtomicU32::new(start),
        }
    }

    pub fn next(&self) -> u32 {
        self.counter.fetch_add(2, Ordering::Relaxed)
    }
}

pub struct MuxMessage {
    pub stream_id: u32,
    pub frame_type: FrameType,
    pub data: Bytes,
}

#[derive(Clone)]
pub struct Muxer {
    control_tx: Sender<MuxMessage>,
    data_tx: Sender<MuxMessage>,
    streams: Arc<DashMap<u32, Sender<Bytes>>>,
    id_gen: Arc<IdGenerator>,
}

impl Muxer {
    pub fn new(
        control_tx: Sender<MuxMessage>,
        data_tx: Sender<MuxMessage>,
        is_client: bool,
    ) -> Self {
        Self {
            control_tx,
            data_tx,
            streams: Arc::new(DashMap::new()),
            id_gen: Arc::new(IdGenerator::new(is_client)),
        }
    }

    pub fn next_id(&self) -> u32 {
        self.id_gen.next()
    }

    // Обычные данные летят в канал с низким приоритетом
    pub async fn send_to_netwrok(&self, message: MuxMessage) -> Result<(), SendError<MuxMessage>> {
        self.data_tx.send(message).await
    }

    // Управляющие фреймы летят в VIP-канал
    pub async fn send_control(
        &self,
        stream_id: u32,
        f_type: FrameType,
        data: Bytes,
    ) -> Result<(), String> {
        self.control_tx
            .send(MuxMessage {
                stream_id,
                frame_type: f_type,
                data,
            })
            .await
            .map_err(|e| e.to_string())
    }

    pub fn register_stream(&self, stream_id: u32, tx: Sender<Bytes>) {
        self.streams.insert(stream_id, tx);
    }

    pub fn remove_stream(&self, stream_id: u32) {
        self.streams.remove(&stream_id);
    }

    pub async fn dispatch_to_local(&self, stream_id: u32, data: Bytes) {
        let tx = self.streams.get(&stream_id).map(|r| r.value().clone());

        if let Some(tx) = tx {
            if let Err(_e) = tx.send(data).await {
                netrunner_logger::warn!(
                    stream_id,
                    "MUXER: [INBOUND_ERR] Local channel closed, dropping packet"
                );
                self.remove_stream(stream_id);
            }
        } else {
            // Уменьшил уровень лога до TRACE, потому что это штатная ситуация при больших буферах!
            netrunner_logger::trace!(
                stream_id,
                len = data.len(),
                "MUXER: [IGNORE] Received data for already closed stream (draining pipe)"
            );
        }
    }
}