use std::num::NonZeroU16;
use std::sync::Arc;
use std::sync::atomic::{AtomicU16, Ordering};

use atuin_common::ansi;

use crate::osc133::{Event, Params, Parser, Zone};

const HISTORY_ID_PARAM: &str = "history_id";
const SESSION_ID_PARAM: &str = "session_id";
const MAX_OUTPUT_CAPTURE_BYTES: usize = 1024 * 1024;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommandCapture {
    pub prompt: String,
    pub command: String,
    pub output: String,
    pub exit_code: Option<i32>,
    pub history_id: Option<String>,
    pub session_id: Option<String>,
    pub output_truncated: bool,
    pub output_observed_bytes: u64,
}

pub type CommandCaptureSink = Box<dyn Fn(CommandCapture) + Send + 'static>;

#[derive(Default)]
struct CaptureBuffers {
    prompt: Vec<u8>,
    command: Vec<u8>,
    output: Vec<u8>,
    output_observed_bytes: u64,
    output_truncated: bool,
    exit_code: Option<i32>,
    history_id: Option<String>,
    session_id: Option<String>,
}

pub struct CommandCaptureTracker {
    parser: Parser,
    zone: Zone,
    buffers: CaptureBuffers,
    cols: Arc<AtomicU16>,
}

impl CommandCaptureTracker {
    pub(crate) fn new(cols: Arc<AtomicU16>) -> Self {
        Self {
            parser: Parser::new(),
            zone: Zone::Unknown,
            buffers: CaptureBuffers::default(),
            cols,
        }
    }

    pub(crate) fn push(&mut self, data: &[u8], mut on_capture: impl FnMut(CommandCapture)) {
        let mut events = Vec::new();
        self.parser.push_located(data, |located| events.push(located));

        let mut start = 0;
        for located in events {
            let marker_start = located.start_offset.min(data.len()).max(start);
            let offset = located.offset.min(data.len());
            self.append(&data[start..marker_start]);
            self.handle_event(located.event, &located.params, &mut on_capture);
            self.zone = located.zone;
            start = offset;
        }

        let append_end = self
            .parser
            .incomplete_osc_sequence_start()
            .map_or(data.len(), |sequence_start| sequence_start.min(data.len()).max(start));
        if start < append_end {
            self.append(&data[start..append_end]);
        }
    }

    fn append(&mut self, data: &[u8]) {
        match self.zone {
            Zone::Prompt => self.buffers.prompt.extend_from_slice(data),
            Zone::Input => self.buffers.command.extend_from_slice(data),
            Zone::Output => self.append_output(data),
            Zone::Unknown => {}
        }
    }

    fn append_output(&mut self, data: &[u8]) {
        self.buffers.output_observed_bytes =
            self.buffers.output_observed_bytes.saturating_add(data.len() as u64);

        if self.buffers.output_truncated {
            return;
        }

        let remaining = MAX_OUTPUT_CAPTURE_BYTES.saturating_sub(self.buffers.output.len());
        let retained = data.len().min(remaining);
        self.buffers.output_truncated = retained < data.len();

        if retained > 0 {
            self.buffers.output.extend_from_slice(&data[..retained]);
        }
    }

    fn handle_event(
        &mut self,
        event: Event,
        params: &Params,
        on_capture: &mut impl FnMut(CommandCapture),
    ) {
        match event {
            Event::PromptStart => {
                if self.zone != Zone::Prompt {
                    self.buffers = CaptureBuffers::default();
                }
            }
            Event::CommandStart | Event::CommandExecuted => {}
            Event::CommandFinished { exit_code } => {
                let Some(history_id) = params.get(HISTORY_ID_PARAM).map(str::to_owned) else {
                    return;
                };

                if exit_code.is_some() || self.buffers.exit_code.is_none() {
                    self.buffers.exit_code = exit_code;
                }
                self.buffers.history_id = Some(history_id);
                self.buffers.session_id = params.get(SESSION_ID_PARAM).map(str::to_owned);

                if let Some(capture) = self.finish_capture() {
                    on_capture(capture);
                }
            }
        }
    }

    fn finish_capture(&mut self) -> Option<CommandCapture> {
        let buffers = std::mem::take(&mut self.buffers);
        // A terminal width of 0 (e.g. before the size is known) falls back to 1.
        let cols = NonZeroU16::new(self.cols.load(Ordering::Relaxed)).unwrap_or(NonZeroU16::MIN);
        let prompt = ansi::to_plain_text(&buffers.prompt, cols);
        let command = ansi::to_plain_text(&buffers.command, cols)
            .trim_matches(|c| c == '\r' || c == '\n')
            .to_string();
        let output = ansi::to_plain_text(&buffers.output, cols);
        let output_truncated = buffers.output_truncated;
        let output_observed_bytes = buffers.output_observed_bytes;
        let exit_code = buffers.exit_code;
        let history_id = buffers.history_id;
        let session_id = buffers.session_id;

        if command.is_empty() && output.is_empty() {
            return None;
        }

        Some(CommandCapture {
            prompt,
            command,
            output,
            exit_code,
            history_id,
            session_id,
            output_truncated,
            output_observed_bytes,
        })
    }
}

#[cfg(test)]
mod tests {
    use rstest::{fixture, rstest};

    use super::*;

    #[fixture]
    fn tracker(#[default(80)] cols: u16) -> CommandCaptureTracker {
        CommandCaptureTracker::new(Arc::new(AtomicU16::new(cols)))
    }

    fn assert_no_terminal_controls(text: &str) {
        assert!(
            !text.chars().any(|ch| ch.is_control() && ch != '\n' && ch != '\t'),
            "text still contains terminal controls: {text:?}"
        );
    }

    #[rstest]
    fn command_text_replays_backspaces(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        let input =
            b"\x1b]133;A\x07$ \x1b]133;B\x07e\x08echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ ";
        tracker.push(input, |capture| captures.push(capture));

        assert_eq!(captures.len(), 1);
        assert_eq!(captures[0].command, "echo hi");
        assert_eq!(captures[0].output, "hi");
        assert_no_terminal_controls(&captures[0].command);
        assert_no_terminal_controls(&captures[0].output);
    }

    #[rstest]
    #[case::complete_command(
        b"\x1b]133;A\x07$ \x1b]133;B\x07echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ ",
        CommandCapture {
            prompt: "$".to_string(),
            command: "echo hi".to_string(),
            output: "hi".to_string(),
            exit_code: Some(0),
            history_id: Some("hist".to_string()),
            session_id: Some("sess".to_string()),
            output_truncated: false,
            output_observed_bytes: 4,
        }
    )]
    #[case::output_only_from_d_marker(
        b"\x1b]133;C\x07line one\r\n\x1b]133;D;0;history_id=018f;session_id=abcd\x07",
        CommandCapture {
            prompt: String::new(),
            command: String::new(),
            output: "line one".to_string(),
            exit_code: Some(0),
            history_id: Some("018f".to_string()),
            session_id: Some("abcd".to_string()),
            output_truncated: false,
            output_observed_bytes: 10,
        }
    )]
    fn captures_full_command(
        mut tracker: CommandCaptureTracker,
        #[case] input: &[u8],
        #[case] expected: CommandCapture,
    ) {
        let mut captures = Vec::new();
        tracker.push(input, |capture| captures.push(capture));
        assert_eq!(captures, vec![expected]);
    }

    #[rstest]
    fn strips_ansi_and_split_markers(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(b"\x1b]133;A\x07\x1b[32m%\x1b[0m ", |_| {});
        tracker.push(b"\x1b]133;B\x07ls\x1b]133;C", |_| {});
        tracker.push(
            b"\x07\x1b[31mfile\x1b[0m\r\n\x1b]133;D;1;history_id=hist;session_id=sess\x07\x1b]133;A\x07% ",
            |capture| {
                captures.push(capture);
            },
        );

        assert_eq!(captures, vec![CommandCapture {
            prompt: "%".to_string(),
            command: "ls".to_string(),
            output: "file".to_string(),
            exit_code: Some(1),
            history_id: Some("hist".to_string()),
            session_id: Some("sess".to_string()),
            output_truncated: false,
            output_observed_bytes: 15,
        }]);
    }

    #[rstest]
    fn duplicate_prompt_start_does_not_reset_prompt_capture(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(
            b"\x1b]133;A\x07$ \x1b]133;A\x07continued \x1b]133;B\x07echo hi\r\n\x1b]133;C\x07hi\r\n\x1b]133;D;0;history_id=hist;session_id=sess\x07\x1b]133;A\x07$ ",
            |capture| captures.push(capture),
        );

        assert_eq!(captures.len(), 1);
        assert_eq!(captures[0].prompt, "$ continued");
        assert_eq!(captures[0].command, "echo hi");
        assert_eq!(captures[0].output, "hi");
    }

    #[rstest]
    fn bare_finish_without_metadata_is_ignored(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(b"\x1b]133;C\x07line one\r\n\x1b]133;D;0\x07", |capture| {
            captures.push(capture);
        });

        tracker.push(b"\x1b]133;A\x07$ ", |capture| captures.push(capture));

        assert!(captures.is_empty());
    }

    #[rstest]
    fn bare_finish_before_metadata_in_same_push_ignored(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(
            b"\x1b]133;C\x07line one\r\n\x1b]133;D;1\x07\x1b]133;D;0;history_id=018f;session_id=abcd\x07",
            |capture| captures.push(capture),
        );

        assert_eq!(captures.len(), 1);
        assert_eq!(captures[0].output, "line one");
        assert_eq!(captures[0].exit_code, Some(0));
        assert_eq!(captures[0].history_id.as_deref(), Some("018f"));
        assert_eq!(captures[0].session_id.as_deref(), Some("abcd"));
    }

    #[rstest]
    fn metadata_arriving_after_bare_finish_across_pushes(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(b"\x1b]133;C\x07line one\r\n\x1b]133;D;0\x07", |capture| {
            captures.push(capture);
        });
        tracker.push(b"\x1b]133;D;0;history_id=018f", |capture| captures.push(capture));

        assert!(captures.is_empty());

        tracker.push(b";session_id=abcd\x07", |capture| captures.push(capture));

        assert_eq!(captures.len(), 1);
        assert_eq!(captures[0].output, "line one");
        assert_eq!(captures[0].exit_code, Some(0));
        assert_eq!(captures[0].history_id.as_deref(), Some("018f"));
        assert_eq!(captures[0].session_id.as_deref(), Some("abcd"));
    }

    #[rstest]
    fn split_finish_marker_is_not_counted_as_output(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(b"\x1b]133;C\x07line one\r\n\x1b]133;D;0;history_id=018f", |capture| {
            captures.push(capture);
        });
        assert!(captures.is_empty());

        tracker.push(b";session_id=abcd\x07", |capture| captures.push(capture));

        assert_eq!(captures.len(), 1);
        assert_eq!(captures[0].output, "line one");
        assert_eq!(captures[0].output_observed_bytes, 10);
    }

    #[rstest]
    fn output_capture_is_capped_and_reports_observed_bytes(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();
        let mut input = b"\x1b]133;C\x07".to_vec();
        input.extend(std::iter::repeat_n(b'x', MAX_OUTPUT_CAPTURE_BYTES + 10));
        input.extend_from_slice(b"\x1b]133;D;0;history_id=big;session_id=session-1\x07");

        tracker.push(&input, |capture| captures.push(capture));

        assert_eq!(captures.len(), 1);
        assert!(captures[0].output_truncated);
        assert_eq!(captures[0].output_observed_bytes, (MAX_OUTPUT_CAPTURE_BYTES + 10) as u64);
    }

    #[rstest]
    fn resets_buffers_between_c_d_only_captures(mut tracker: CommandCaptureTracker) {
        let mut captures = Vec::new();

        tracker.push(
            b"\x1b]133;C\x07first\r\n\x1b]133;D;0;history_id=one\x07\x1b]133;C\x07second\r\n\x1b]133;D;1;history_id=two\x07",
            |capture| captures.push(capture),
        );

        assert_eq!(captures.len(), 2);
        assert_eq!(captures[0].output, "first");
        assert_eq!(captures[0].history_id.as_deref(), Some("one"));
        assert_eq!(captures[1].output, "second");
        assert_eq!(captures[1].history_id.as_deref(), Some("two"));
    }
}
