use chrono::{SecondsFormat, Utc};
use feldera_observability::json_logging::use_json_log_format;
use serde_json::json;
use std::collections::VecDeque;
use std::time::Duration;
use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError};
use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
use tokio::time::interval;
use tokio::{select, spawn};
use tracing::{debug, error, warn, Level};

// Logs buffer size limit constants.
const LOGS_BUFFER_LIMIT_BYTE: usize = 1_000_000; // 1 MB
const LOGS_BUFFER_LIMIT_NUM_LINES: usize = 50_000; // 50K lines

// Timeout to send to the pipeline logs.
const SEND_LOG_MESSAGE_TIMEOUT: Duration = Duration::from_millis(100);

// Number of time to try sending the logs message before giving up.
const SEND_LOG_MESSAGE_TRIES: u64 = 100;

#[derive(Clone)]
pub enum LogMessage {
    Pipeline {
        line: String,
    },
    ControlPlane {
        target: &'static str,
        service: &'static str,
        pipeline_name: String,
        pipeline_id: String,
        level: Level,
        line: String,
        timestamp: String,
    },
}

impl LogMessage {
    /// Constructs a log message for a line generated by the pipeline itself.
    /// The pipeline log line already contains a timestamp and level.
    pub fn new_from_pipeline(line: &str) -> LogMessage {
        LogMessage::Pipeline {
            line: line.to_string(),
        }
    }

    /// Constructs a log message for a line generated by the runner (automaton or executor).
    pub fn new_from_control_plane(
        target: &'static str,
        service: &'static str,
        pipeline_name: impl Into<String>,
        pipeline_id: impl Into<String>,
        level: Level,
        line: &str,
    ) -> LogMessage {
        LogMessage::ControlPlane {
            target,
            service,
            pipeline_name: pipeline_name.into(),
            pipeline_id: pipeline_id.into(),
            level,
            line: line.to_string(),
            timestamp: Utc::now().to_rfc3339_opts(SecondsFormat::Micros, true),
        }
    }
}

/// Sets up a thread which listens to follow requests and new incoming log lines.
/// New followers are caught up and existing followers receive new lines as they come in.
/// Returns a termination sender and the corresponding join handle.
#[allow(clippy::type_complexity)]
pub fn start_thread_pipeline_logs(
    pipeline_id: impl Into<String>,
    pipeline_name: impl Into<String>,
    mut follow_request_receiver: mpsc::Receiver<mpsc::Sender<String>>,
    mut logs_receiver: mpsc::Receiver<LogMessage>,
) -> (oneshot::Sender<()>, JoinHandle<()>) {
    let pipeline_id = pipeline_id.into();
    let pipeline_name = pipeline_name.into();
    let (terminate_sender, mut terminate_receiver) = oneshot::channel::<()>();
    let join_handle = spawn(async move {
        // Buffer with the latest lines
        let mut logs = LogsBuffer::new(LOGS_BUFFER_LIMIT_BYTE, LOGS_BUFFER_LIMIT_NUM_LINES);

        // First line
        logs.append(format_log_line(&LogMessage::new_from_control_plane(
            module_path!(),
            "runner",
            pipeline_name.clone(),
            pipeline_id.clone(),
            Level::INFO,
            "Fresh start of pipeline logs",
        )));

        // Followers interested in receiving the logs
        let mut log_followers: Vec<mpsc::Sender<String>> = Vec::new();

        // Normally dead followers are dropped when we receive new message.
        // This just helps in cleanup without the need for new message.
        // Hence, we can just do it if idle for 5 mins, as in ideal case,
        // there would be new messages before it.
        let mut idle_interval_for_cleanup = interval(Duration::from_secs(300));

        loop {
            select! {
                // Termination request
                _ = &mut terminate_receiver => {
                    break;
                }

                // Follow request
                follower = follow_request_receiver.recv() => {
                    if let Some(follower) = follower {
                        catch_up_and_add_follower(
                            &pipeline_id,
                            &pipeline_name,
                            &mut logs,
                            &mut log_followers,
                            follower,
                        )
                        .await;
                    } else {
                        // The follow request sender has been dropped, which occurs when the pipeline is deleted.
                        // In this case, the logs thread is also terminated.
                        break;
                    }
                }

                // Log message
                message = logs_receiver.recv() => {
                    match message {
                        Some(message) => {
                            let line = format_log_line(&message);
                            process_log_line_with_followers(
                                &mut logs,
                                &mut log_followers,
                                line
                            ).await;
                        },
                        None => {
                            // All logs senders have been dropped, which can occur when the pipeline is deleted.
                            // In this case, the logs thread is also terminated.
                            break;
                        }
                    }
                }

                _ = idle_interval_for_cleanup.tick() => {
                    // drop the dead followers
                    log_followers.retain(|follower| !follower.is_closed());
                }
            }
        }
    });
    (terminate_sender, join_handle)
}

/// Catches up the follower by sending all the buffered logs to it.
/// Afterward, adds it to the list of known followers if there was
/// no error during sending the catch-up.
async fn catch_up_and_add_follower(
    pipeline_id: &str,
    pipeline_name: &str,
    logs: &mut LogsBuffer,
    log_followers: &mut Vec<mpsc::Sender<String>>,
    new_follower: mpsc::Sender<String>,
) {
    // Catch up the new follower if there are any lines to catch up on
    let mut failed = false;

    // First line mentions the number of discarded lines due to the circular buffer
    if logs.num_discarded_lines() > 0 {
        let first_line = format!("... {} prior log lines were discarded due to buffer constraints and are thus not shown.", logs.num_discarded_lines());
        // Tag as control-plane metadata so the notice is formatted (text or JSON) consistently with other runner messages.
        let formatted_notice = format_log_line(&LogMessage::new_from_control_plane(
            module_path!(),
            "runner",
            pipeline_name.to_string(),
            pipeline_id.to_string(),
            Level::WARN,
            &first_line,
        ));
        if let Err(e) = new_follower.try_send(formatted_notice) {
            match e {
                TrySendError::Full(_) => {
                    error!("Unable to catch up new follower because buffer is full, the follower will be dropped");
                }
                TrySendError::Closed(_) => {}
            }
            failed = true;
        }
    }

    // Feed all the lines stored in the circular buffer
    if !failed {
        for line in logs.lines() {
            if let Err(e) = new_follower.try_send(line.clone()) {
                match e {
                    TrySendError::Full(_) => {
                        error!("Unable to catch up new follower because buffer is full, the follower will be dropped")
                    }
                    TrySendError::Closed(_) => {}
                }
                failed = true;
                break;
            }
        }
    }

    // Any failure in sending results in the follower Sender not being added
    // to the list, and thus going out of scope and being dropped.
    // The Receiver in that case will be notified no Sender exists anymore.
    if !failed {
        log_followers.push(new_follower);
    }
}

/// Process a new log line by adding it to the lines buffer and
/// sending it out to all followers. Any followers that exhibit
/// a send error are removed.
async fn process_log_line_with_followers(
    logs: &mut LogsBuffer,
    log_followers: &mut Vec<mpsc::Sender<String>>,
    line: String,
) {
    // Add copy of line to buffer
    logs.append(line.clone());

    // Send to all followers the new line
    let mut retain_indexes = vec![];
    for (idx, follower) in log_followers.iter().enumerate() {
        match follower.try_send(line.clone()) {
            Ok(()) => {
                retain_indexes.push(idx);
            }
            Err(e) => match e {
                TrySendError::Full(_) => {
                    // If the follower is unable to keep up, it will be removed.
                    // There exists a buffer to give a follower the chance to catch up.
                    // However, if the limit of the buffer is reached and thus unable to send new,
                    // the log follower will be removed to prevent it from slowing down the rest.
                    error!("Unable to send log line to follower because buffer is full: the follower will be removed")
                }
                TrySendError::Closed(_) => {}
            },
        }
    }

    // Only keep the followers to who we were able to send the new line
    let mut idx: usize = 0;
    log_followers.retain(|_follower| {
        let keep = retain_indexes.contains(&idx);
        idx += 1;
        keep
    });

    // Any Senders that were not retained will go out of scope, which
    // results in them being dropped and the Receiver being notified
    // no Sender exists anymore.
}

fn format_log_line(message: &LogMessage) -> String {
    if !use_json_log_format() {
        return match message {
            LogMessage::Pipeline { line } => line.clone(),
            LogMessage::ControlPlane {
                target,
                service: _,
                pipeline_name,
                pipeline_id,
                level,
                line,
                timestamp,
            } => {
                let display_name = if pipeline_name.is_empty() {
                    "N/A"
                } else {
                    pipeline_name.as_str()
                };
                let mut output = format!("[manager] {timestamp} {level:>5} {target}: {line}");
                if !pipeline_id.is_empty() {
                    output.push_str(&format!(" pipeline-id={pipeline_id:?}"));
                }
                if !display_name.is_empty() {
                    output.push_str(&format!(" pipeline-name={display_name:?}"));
                }
                output
            }
        };
    }

    match message {
        // Pipeline lines are already formatted, forward as-is
        LogMessage::Pipeline { line } => line.clone(),
        LogMessage::ControlPlane {
            target,
            service,
            pipeline_name,
            pipeline_id,
            level,
            line,
            timestamp,
        } => json!({
            "timestamp": timestamp,
            "level": level.as_str(),
            "target": target,
            "feldera-service": service,
            "pipeline-name": pipeline_name,
            "pipeline-id": pipeline_id,
            "fields": { "line": line },
        })
        .to_string(),
    }
}

/// Wrapper around the logs sender channel, which gracefully handles inability to send a message.
#[derive(Clone)]
pub struct LogsSender {
    sender: mpsc::Sender<LogMessage>,
}

impl LogsSender {
    pub fn new(sender: mpsc::Sender<LogMessage>) -> Self {
        Self { sender }
    }

    pub async fn send(&mut self, mut message: LogMessage) {
        // This will momentarily block when the receiver buffer is full. This should generally not
        // happen as the receiving thread continuously listens for new log messages and puts them
        // into a circular buffer. It retries a set amount of times with a timeout inbetween before
        // giving up.
        for i in 1..=SEND_LOG_MESSAGE_TRIES {
            message = match self
                .sender
                .send_timeout(message, SEND_LOG_MESSAGE_TIMEOUT)
                .await
            {
                Ok(()) => {
                    // Successfully sent
                    return;
                }
                Err(e) => match e {
                    SendTimeoutError::Timeout(unsent_message) => {
                        warn!(
                            "Unable to send logs message because receiver buffer is full -- trying again in {}ms (attempt {} / {})",
                            SEND_LOG_MESSAGE_TIMEOUT.as_millis(), i, SEND_LOG_MESSAGE_TRIES
                        );
                        unsent_message
                    }
                    SendTimeoutError::Closed(_) => {
                        debug!("Unable to send logs message because receiver is closed -- this can happen when the pipeline is deleted");
                        return;
                    }
                },
            }
        }
        error!(
            "Unable to send logs message after attempting {SEND_LOG_MESSAGE_TRIES} times -- receiver buffer is consistently full: message (byte length: {}) is dropped",
            format_log_line(&message).len()
        )
    }
}

/// The LogsBuffer maintains internally a circular buffer of Strings whose
/// size in byte and number of elements does not exceed the limits.
/// When appending new log lines (Strings) to the buffer, the limits are
/// enforced by discarding existing lines if the limits would be exceeded.
pub struct LogsBuffer {
    /// Buffer size limit in byte.
    size_limit_byte: usize,
    /// Buffer size limit in number of lines.
    size_limit_num_lines: usize,
    /// The lines buffer.
    buffer: VecDeque<String>,
    /// Current lines buffer size.
    size_byte: usize,
    /// Number of lines that have been discarded to enforce size limit.
    num_discarded_lines: usize,
}

impl LogsBuffer {
    /// Construct a new logs buffer.
    pub fn new(size_limit_byte: usize, size_limit_num_lines: usize) -> Self {
        Self {
            size_limit_byte,
            size_limit_num_lines,
            buffer: VecDeque::new(),
            size_byte: 0,
            num_discarded_lines: 0,
        }
    }

    /// Append a new line to the buffer.
    /// - If the new line exceeds the buffer size limit by itself, all the lines in the buffer and
    ///   the new line are discarded, leaving an empty buffer.
    /// - Otherwise, lines are removed from the buffer until the new line will fit. Once there
    ///   is sufficient space, the line is added to the buffer.
    pub fn append(&mut self, line: String) {
        if line.len() > self.size_limit_byte {
            self.num_discarded_lines += self.buffer.len() + 1;
            self.buffer.clear();
            self.size_byte = 0;
        } else {
            // Ensure size in byte is not exceeded
            while self.size_byte + line.len() > self.size_limit_byte {
                let popped_line = self
                    .buffer
                    .pop_front()
                    .expect("Cannot remove log line even though size is non-zero");
                self.size_byte -= popped_line.len();
                self.num_discarded_lines += 1;
            }

            // Ensure size in number of lines is not exceeded
            if self.size_limit_num_lines > 0 {
                while self.buffer.len() + 1 > self.size_limit_num_lines {
                    let popped_line = self
                        .buffer
                        .pop_front()
                        .expect("Cannot remove log line even though length is non-zero");
                    self.size_byte -= popped_line.len();
                    self.num_discarded_lines += 1;
                }
                self.size_byte += line.len();
                self.buffer.push_back(line);
            }
        }
    }

    /// Retrieves the lines in the buffer.
    pub fn lines(&self) -> &VecDeque<String> {
        &self.buffer
    }

    /// Retrieves the number of lines in the buffer.
    pub fn num_lines(&self) -> usize {
        self.buffer.len()
    }

    /// Retrieves the number of lines discarded due to buffer limit enforcement.
    pub fn num_discarded_lines(&self) -> usize {
        self.num_discarded_lines
    }

    /// Retrieves the total buffer size.
    pub fn size_byte(&self) -> usize {
        self.size_byte
    }

    /// Retrieves the buffer size limit in byte.
    pub fn size_limit_byte(&self) -> usize {
        self.size_limit_byte
    }

    /// Retrieves the buffer size limit in number of lines.
    pub fn size_limit_num_lines(&self) -> usize {
        self.size_limit_num_lines
    }
}

#[cfg(test)]
mod test {
    use super::LogsBuffer;
    use std::collections::VecDeque;

    #[test]
    #[rustfmt::skip] // Skip formatting to keep the test cases readable
    fn logs_buffer_variety() {
        let test_cases: Vec<(usize, usize, Vec<&str>, Vec<&str>)> = vec![
            // Potentially exceed number of byte
            (0, 1000, vec!["a"], vec![]),
            (0, 1000, vec!["a"], vec![]),
            (0, 1000, vec!["a", "b"], vec![]),
            (1, 1000, vec!["a"], vec!["a"]),
            (1, 1000, vec!["a", "b"], vec!["b"]),
            (1, 1000, vec!["a", "b", "c"], vec!["c"]),
            (2, 1000, vec!["a", "b", "c"], vec!["b", "c"]),
            (2, 1000, vec!["a", "b", "c", "d"], vec!["c", "d"]),
            // Potentially exceed number of lines
            (1000, 0, vec![], vec![]),
            (1000, 0, vec!["a"], vec![]),
            (1000, 0, vec!["a", "b"], vec![]),
            (1000, 1, vec!["a"], vec!["a"]),
            (1000, 1, vec!["a", "b"], vec!["b"]),
            (1000, 1, vec!["a", "b", "c"], vec!["c"]),
            (1000, 2, vec!["a", "b", "c"], vec!["b", "c"]),
            (1000, 2, vec!["a", "b", "c", "d"], vec!["c", "d"]),
            // Empty lines exceed number of lines
            (1000, 0, vec![""], vec![]),
            (1000, 1, vec![""], vec![""]),
            (1000, 1, vec!["", ""], vec![""]),
            (1000, 5, vec!["", "", "", "", "", ""], vec!["", "", "", "", ""]),
            // Exceed both potentially
            (0, 0, vec!["a"], vec![]),
            (0, 0, vec!["a", "b"], vec![]),
            (1, 1, vec!["a"], vec!["a"]),
            (1, 1, vec!["a", "b"], vec!["b"]),
            (1, 1, vec!["a", "b", "c"], vec!["c"]),
            (2, 2, vec!["a", "b", "c"], vec!["b", "c"]),
            (2, 2, vec!["a", "b", "c", "d"], vec!["c", "d"]),
            // Others
            (5, 10, vec!["abc", "def"], vec!["def"]),
            (5, 2, vec!["abc", "def"], vec!["def"]),
            (6, 2, vec!["abc", "def"], vec!["abc", "def"]),
        ];

        // Run the test cases
        for (limit_byte, limit_num_lines, input, output) in test_cases {
            let mut buffer = LogsBuffer::new(limit_byte, limit_num_lines);
            for s in &input {
                buffer.append(s.to_string().clone());
            }
            let mut expected = VecDeque::new();
            for s in &output {
                expected.push_back(s.to_string());
            }
            assert_eq!(buffer.lines(), &expected, "Failed for test case (lb={}, ln={}, i={:?}) -> o={:?}", limit_byte, limit_num_lines, input, output);
        }
    }

    #[test]
    fn logs_buffer_normal() {
        // Buffer with 20 byte and 5 lines limit
        let mut buffer = LogsBuffer::new(20, 5);
        assert_eq!(buffer.lines(), &VecDeque::from([]));
        assert_eq!(buffer.num_lines(), 0);
        assert_eq!(buffer.num_discarded_lines(), 0);
        assert_eq!(buffer.size_byte(), 0);
        assert_eq!(buffer.size_limit_byte(), 20);
        assert_eq!(buffer.size_limit_num_lines(), 5);

        // Exceed the byte limit
        buffer.append("abcde".to_string());
        buffer.append("fghij".to_string());
        buffer.append("klmno".to_string());
        buffer.append("pqrst".to_string());
        assert_eq!(buffer.num_discarded_lines(), 0);
        buffer.append("uvwxyz1".to_string());
        assert_eq!(
            buffer.lines(),
            &VecDeque::from([
                "klmno".to_string(),
                "pqrst".to_string(),
                "uvwxyz1".to_string()
            ])
        );
        assert_eq!(buffer.num_lines(), 3);
        assert_eq!(buffer.num_discarded_lines(), 2);
        assert_eq!(buffer.size_byte(), 17);
        assert_eq!(buffer.size_limit_byte(), 20);
        assert_eq!(buffer.size_limit_num_lines(), 5);

        // Exceed the number of lines limit
        buffer.append("2".to_string());
        buffer.append("3".to_string());
        assert_eq!(buffer.num_discarded_lines(), 2);
        buffer.append("4".to_string());
        assert_eq!(
            buffer.lines(),
            &VecDeque::from([
                "pqrst".to_string(),
                "uvwxyz1".to_string(),
                "2".to_string(),
                "3".to_string(),
                "4".to_string()
            ])
        );
        assert_eq!(buffer.num_lines(), 5);
        assert_eq!(buffer.num_discarded_lines(), 3);
        assert_eq!(buffer.size_byte(), 15);
        assert_eq!(buffer.size_limit_byte(), 20);
        assert_eq!(buffer.size_limit_num_lines(), 5);

        // Exceed the number of bytes with a string larger than can fit in buffer
        buffer.append("aaaaabbbbbcccccddddde".to_string());
        assert_eq!(buffer.num_lines(), 0);
        assert_eq!(buffer.num_discarded_lines(), 9);
        assert_eq!(buffer.size_byte(), 0);
        assert_eq!(buffer.size_limit_byte(), 20);
        assert_eq!(buffer.size_limit_num_lines(), 5);
    }
}
