use crate::{
    smp::core::smp_get_processor_id,
    tracepoint::{TraceEntry, TracePointMap},
};
use alloc::{format, string::String, vec::Vec};

pub trait TracePipeOps {
    /// Returns the first event in the trace pipe buffer without removing it.
    fn peek(&self) -> Option<&Vec<u8>>;

    /// Remove and return the first event in the trace pipe buffer.
    fn pop(&mut self) -> Option<Vec<u8>>;

    /// Whether the trace pipe buffer is empty.
    fn is_empty(&self) -> bool;
}

/// A raw trace pipe buffer that stores trace events as byte vectors.
pub struct TracePipeRaw {
    max_record: usize,
    event_buf: Vec<Vec<u8>>,
}

impl TracePipeRaw {
    pub const fn new(max_record: usize) -> Self {
        Self {
            max_record,
            event_buf: Vec::new(),
        }
    }

    /// Set the maximum number of records to keep in the trace pipe buffer.
    ///
    /// If the current number of records exceeds this limit, the oldest records will be removed.
    #[allow(unused)]
    pub fn set_max_record(&mut self, max_record: usize) {
        self.max_record = max_record;
        if self.event_buf.len() > max_record {
            self.event_buf.truncate(max_record); // Keep only the latest records
        }
    }

    /// Push a new event into the trace pipe buffer.
    pub fn push_event(&mut self, event: Vec<u8>) {
        if self.event_buf.len() >= self.max_record {
            self.event_buf.remove(0); // Remove the oldest record
        }
        self.event_buf.push(event);
    }

    /// Clear the trace pipe buffer.
    pub fn clear(&mut self) {
        self.event_buf.clear();
    }

    /// Create a snapshot of the current state of the trace pipe buffer.
    pub fn snapshot(&self) -> TracePipeSnapshot {
        TracePipeSnapshot::new(self.event_buf.clone())
    }
}

impl TracePipeOps for TracePipeRaw {
    fn peek(&self) -> Option<&Vec<u8>> {
        self.event_buf.first()
    }

    fn pop(&mut self) -> Option<Vec<u8>> {
        if self.event_buf.is_empty() {
            None
        } else {
            Some(self.event_buf.remove(0))
        }
    }

    fn is_empty(&self) -> bool {
        self.event_buf.is_empty()
    }
}

#[derive(Debug, Clone)]
pub struct TracePipeSnapshot(Vec<Vec<u8>>);

impl TracePipeSnapshot {
    pub fn new(event_buf: Vec<Vec<u8>>) -> Self {
        Self(event_buf)
    }

    /// The formatted string representation to be used as a header for the trace pipe output.
    pub fn default_fmt_str(&self) -> String {
        let show = "#
#
#                                _-----=> irqs-off/BH-disabled
#                               / _----=> need-resched
#                              | / _---=> hardirq/softirq
#                              || / _--=> preempt-depth
#                              ||| / _-=> migrate-disable
#                              |||| /     delay
#           TASK-PID     CPU#  |||||  TIMESTAMP  FUNCTION
#              | |         |   |||||     |         |
";
        format!(
            "# tracer: nop\n#\n# entries-in-buffer/entries-written: {}/{}   #P:32\n{}",
            self.0.len(),
            self.0.len(),
            show
        )
    }
}

impl TracePipeOps for TracePipeSnapshot {
    fn peek(&self) -> Option<&Vec<u8>> {
        self.0.first()
    }

    fn pop(&mut self) -> Option<Vec<u8>> {
        if self.0.is_empty() {
            None
        } else {
            Some(self.0.remove(0))
        }
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

/// A cache for storing command line arguments for each trace point.
///
/// See https://www.kernel.org/doc/Documentation/trace/ftrace.txt
pub struct TraceCmdLineCache {
    cmdline: Vec<(u32, [u8; 16])>,
    max_record: usize,
}

impl TraceCmdLineCache {
    pub const fn new(max_record: usize) -> Self {
        Self {
            cmdline: Vec::new(),
            max_record,
        }
    }

    /// Insert a command line argument for a trace point.
    ///
    /// If the command line exceeds 16 bytes, it will be truncated.
    /// If the cache exceeds the maximum record limit, the oldest entry will be removed.
    pub fn insert(&mut self, id: u32, cmdline: String) {
        if self.max_record == 0 {
            return;
        }
        if self.cmdline.len() >= self.max_record {
            // Remove the oldest entry if we exceed the max record limit
            self.cmdline.remove(0);
        }
        let mut cmdline_bytes = [0u8; 16];
        let mut len = cmdline.len().min(15);
        while len > 0 && !cmdline.is_char_boundary(len) {
            len -= 1;
        }
        cmdline_bytes[..len].copy_from_slice(&cmdline.as_bytes()[..len]);
        self.cmdline.push((id, cmdline_bytes));
    }

    /// Get the command line argument for a trace point.
    pub fn get(&self, id: u32) -> Option<&str> {
        self.cmdline.iter().find_map(|(key, value)| {
            if *key == id {
                let len = value
                    .iter()
                    .position(|byte| *byte == 0)
                    .unwrap_or(value.len());
                core::str::from_utf8(&value[..len]).ok()
            } else {
                None
            }
        })
    }

    /// Set the maximum length for command line arguments.
    pub fn set_max_record(&mut self, max_len: usize) {
        self.max_record = max_len;
        if self.cmdline.len() > max_len {
            let remove = self.cmdline.len() - max_len;
            self.cmdline.drain(..remove);
        }
    }

    /// Get the maximum number of records in the cache.
    pub fn max_record(&self) -> usize {
        self.max_record
    }

    /// Create a snapshot of the current state of the command line cache.
    pub fn snapshot(&self) -> TraceCmdLineCacheSnapshot {
        TraceCmdLineCacheSnapshot::new(self.cmdline.clone())
    }
}

#[derive(Debug, Clone)]
pub struct TraceCmdLineCacheSnapshot(Vec<(u32, [u8; 16])>);
impl TraceCmdLineCacheSnapshot {
    pub fn new(cmdline: Vec<(u32, [u8; 16])>) -> Self {
        Self(cmdline)
    }

    /// Return the first command line entry in the cache.
    pub fn peek(&self) -> Option<&(u32, [u8; 16])> {
        self.0.first()
    }

    /// Remove and return the first command line entry in the cache.
    pub fn pop(&mut self) -> Option<(u32, [u8; 16])> {
        if self.0.is_empty() {
            None
        } else {
            Some(self.0.remove(0))
        }
    }
}

#[cfg(test)]
mod cmdline_cache_tests {
    use super::TraceCmdLineCache;
    use alloc::string::ToString;

    #[test]
    fn truncates_at_utf8_boundary_and_keeps_nul_space() {
        let mut cache = TraceCmdLineCache::new(4);
        cache.insert(1, "abcdefghijklmnop".to_string());
        assert_eq!(cache.get(1), Some("abcdefghijklmno"));

        cache.insert(2, "abcdefghijklmn界".to_string());
        assert_eq!(cache.get(2), Some("abcdefghijklmn"));
    }

    #[test]
    fn invalid_cached_bytes_do_not_panic() {
        let mut cache = TraceCmdLineCache::new(1);
        let mut invalid = [0u8; 16];
        invalid[0] = 0xff;
        cache.cmdline.push((1, invalid));
        assert_eq!(cache.get(1), None);
    }

    #[test]
    fn zero_capacity_and_shrink_keep_cache_invariants() {
        let mut disabled = TraceCmdLineCache::new(0);
        disabled.insert(1, "ignored".to_string());
        assert_eq!(disabled.get(1), None);

        let mut cache = TraceCmdLineCache::new(4);
        for id in 1..=4 {
            cache.insert(id, id.to_string());
        }
        cache.set_max_record(2);
        assert_eq!(cache.get(1), None);
        assert_eq!(cache.get(2), None);
        assert_eq!(cache.get(3), Some("3"));
        assert_eq!(cache.get(4), Some("4"));
    }
}

pub struct TraceEntryParser;

impl TraceEntryParser {
    /// Parse the trace entry and return a formatted string.
    pub fn parse(
        tracepoint_map: &TracePointMap,
        cmdline_cache: &TraceCmdLineCache,
        entry: &[u8],
    ) -> String {
        let trace_entry = unsafe { &*(entry.as_ptr() as *const TraceEntry) };
        let id = trace_entry.type_ as u32;
        let tracepoint = tracepoint_map.get(&id).expect("TracePoint not found");
        let fmt_func = tracepoint.fmt_func();
        let offset = core::mem::size_of::<TraceEntry>();
        let str = fmt_func(&entry[offset..]);

        let time = crate::time::Instant::now().total_micros() * 1000; // Convert to nanoseconds
        let cpu_id = smp_get_processor_id().data();

        // Copy the packed field to a local variable to avoid unaligned reference
        let pid = trace_entry.pid;
        let pname = cmdline_cache.get(pid as u32).unwrap_or("<...>");

        let secs = time / 1_000_000_000;
        let usec_rem = time % 1_000_000_000 / 1000;

        format!(
            "{:>16}-{:<7} [{:03}] {} {:5}.{:06}: {}({})\n",
            pname,
            pid,
            cpu_id,
            trace_entry.trace_print_lat_fmt(),
            secs,
            usec_rem,
            tracepoint.name(),
            str
        )
    }
}
