use std::collections::HashMap;
use std::sync::Mutex;

use crate::api::properties::ChannelBaseOn;
use crate::api::runtime::{ChannelKey, TaskId};
use crate::channel::{named_channel_with_base, ElementReceiver, ElementSender};
use crate::metrics::Tag;

lazy_static! {
    static ref MEMORY_CHANNELS: Mutex<HashMap<TaskId, (ElementSender, ElementReceiver)>> =
        Mutex::new(HashMap::new());
}

pub(crate) fn publish(
    source_task_id: &TaskId,
    target_task_ids: &Vec<TaskId>,
    channel_size: usize,
    channel_base_on: ChannelBaseOn,
) -> Vec<(ChannelKey, ElementSender)> {
    let mut senders = Vec::new();
    for target_task_id in target_task_ids {
        let channel_key = ChannelKey {
            source_task_id: source_task_id.clone(),
            target_task_id: target_task_id.clone(),
        };
        let sender = get(*target_task_id, channel_size, channel_base_on).0;
        senders.push((channel_key, sender));
    }
    senders
}

pub(crate) fn subscribe(
    source_task_ids: &Vec<TaskId>,
    target_task_id: &TaskId,
    channel_size: usize,
    channel_base_on: ChannelBaseOn,
) -> ElementReceiver {
    if source_task_ids.len() == 0 {
        panic!("source TaskId not found");
    }

    get(*target_task_id, channel_size, channel_base_on).1
}

pub(crate) fn get(
    target_task_id: TaskId,
    channel_size: usize,
    channel_base_on: ChannelBaseOn,
) -> (ElementSender, ElementReceiver) {
    let memory_channels: &Mutex<HashMap<TaskId, (ElementSender, ElementReceiver)>> =
        &*MEMORY_CHANNELS;
    let mut guard = memory_channels.lock().unwrap();
    let (sender, receiver) = guard.entry(target_task_id).or_insert_with(|| {
        named_channel_with_base(
            "Memory_PubSub",
            vec![
                Tag::from(("target_job_id", target_task_id.job_id.0)),
                Tag::from(("target_task_number", target_task_id.task_number)),
            ],
            channel_size,
            channel_base_on,
        )
    });
    (sender.clone(), receiver.clone())
}
