use super::{FilterOption, FilterPlugin};
use crate::common::PIPESIZE;
use crate::services::filter_service::register_filter_plugin;
use crate::tokenizer::{TokenEncoding, Tokenizer, get_tokenizer};
use std::io::{Read, Result, Write};

// ---------------------------------------------------------------------------
// head_tokens
// ---------------------------------------------------------------------------

/// A filter that outputs only the first N tokens of the input stream.
///
/// Streams bytes directly until the token limit is reached. When the limit
/// falls mid-chunk, uses `split_by_token_iter` to find the exact byte boundary
/// without allocating token strings beyond what is needed.
pub struct HeadTokensFilter {
    pub remaining: usize,
    pub tokenizer: Tokenizer,
    pub encoding: TokenEncoding,
}

impl HeadTokensFilter {
    pub fn new(count: usize) -> Self {
        let encoding = TokenEncoding::default();
        Self {
            remaining: count,
            tokenizer: get_tokenizer(encoding).clone(),
            encoding,
        }
    }
}

impl FilterPlugin for HeadTokensFilter {
    fn filter(&mut self, reader: &mut dyn Read, writer: &mut dyn Write) -> Result<()> {
        if self.remaining == 0 {
            return Ok(());
        }

        let tokenizer = &self.tokenizer;
        let mut buffer = vec![0u8; PIPESIZE];
        let mut total_tokens = 0usize;

        loop {
            let n = reader.read(&mut buffer)?;
            if n == 0 {
                break;
            }

            let chunk = &buffer[..n];
            let text = String::from_utf8_lossy(chunk);
            let chunk_tokens = tokenizer.count(&text);

            if total_tokens + chunk_tokens <= self.remaining {
                // Entire chunk fits — write it directly
                writer.write_all(chunk)?;
                total_tokens += chunk_tokens;
                if total_tokens >= self.remaining {
                    break;
                }
            } else {
                // Cutoff is within this chunk — use iterator to find exact
                // boundary without allocating all token strings
                let tokens_to_write = self.remaining - total_tokens;
                let mut byte_pos = 0usize;
                for token_str in tokenizer.split_by_token_iter(&text).take(tokens_to_write) {
                    byte_pos += token_str
                        .map_err(|e| std::io::Error::other(e.to_string()))?
                        .len();
                }
                let write_len = map_lossy_pos_to_bytes(chunk, &text, byte_pos);
                writer.write_all(&chunk[..write_len])?;
                break;
            }
        }
        Ok(())
    }

    fn clone_box(&self) -> Box<dyn FilterPlugin> {
        Box::new(Self {
            remaining: self.remaining,
            tokenizer: get_tokenizer(self.encoding).clone(),
            encoding: self.encoding,
        })
    }

    fn options(&self) -> Vec<FilterOption> {
        vec![
            FilterOption {
                name: "count".to_string(),
                default: None,
                required: true,
            },
            FilterOption {
                name: "encoding".to_string(),
                default: Some(serde_json::Value::String("cl100k_base".to_string())),
                required: false,
            },
        ]
    }

    fn description(&self) -> &str {
        "Read the first N LLM tokens"
    }
}

// ---------------------------------------------------------------------------
// skip_tokens
// ---------------------------------------------------------------------------

/// A filter that skips the first N tokens of the input stream and outputs the rest.
pub struct SkipTokensFilter {
    pub remaining: usize,
    pub tokenizer: Tokenizer,
    pub encoding: TokenEncoding,
}

impl SkipTokensFilter {
    pub fn new(count: usize) -> Self {
        let encoding = TokenEncoding::default();
        Self {
            remaining: count,
            tokenizer: get_tokenizer(encoding).clone(),
            encoding,
        }
    }
}

impl FilterPlugin for SkipTokensFilter {
    fn filter(&mut self, reader: &mut dyn Read, writer: &mut dyn Write) -> Result<()> {
        if self.remaining == 0 {
            return std::io::copy(reader, writer).map(|_| ());
        }

        let tokenizer = &self.tokenizer;
        let mut buffer = vec![0u8; PIPESIZE];
        let mut total_tokens = 0usize;
        let mut done_skipping = false;

        loop {
            let n = reader.read(&mut buffer)?;
            if n == 0 {
                break;
            }

            if done_skipping {
                writer.write_all(&buffer[..n])?;
                continue;
            }

            let chunk = &buffer[..n];
            let text = String::from_utf8_lossy(chunk);
            let chunk_tokens = tokenizer.count(&text);

            if total_tokens + chunk_tokens <= self.remaining {
                // Entire chunk is skipped
                total_tokens += chunk_tokens;
                if total_tokens >= self.remaining {
                    done_skipping = true;
                }
            } else {
                // Cutoff is within this chunk — use iterator to skip past
                // the boundary without allocating all token strings
                let tokens_to_skip = self.remaining - total_tokens;
                let mut byte_pos = 0usize;
                for token_str in tokenizer.split_by_token_iter(&text).take(tokens_to_skip) {
                    byte_pos += token_str
                        .map_err(|e| std::io::Error::other(e.to_string()))?
                        .len();
                }
                let skip_len = map_lossy_pos_to_bytes(chunk, &text, byte_pos);
                if skip_len < n {
                    writer.write_all(&chunk[skip_len..])?;
                }
                done_skipping = true;
            }
        }
        Ok(())
    }

    fn clone_box(&self) -> Box<dyn FilterPlugin> {
        Box::new(Self {
            remaining: self.remaining,
            tokenizer: get_tokenizer(self.encoding).clone(),
            encoding: self.encoding,
        })
    }

    fn options(&self) -> Vec<FilterOption> {
        vec![
            FilterOption {
                name: "count".to_string(),
                default: None,
                required: true,
            },
            FilterOption {
                name: "encoding".to_string(),
                default: Some(serde_json::Value::String("cl100k_base".to_string())),
                required: false,
            },
        ]
    }

    fn description(&self) -> &str {
        "Skip the first N LLM tokens"
    }
}

// ---------------------------------------------------------------------------
// tail_tokens
// ---------------------------------------------------------------------------

/// A filter that outputs only the last N tokens of the input stream.
///
/// Buffers all bytes from the stream, then at finalize tokenizes the
/// content and writes only the last N tokens.
pub struct TailTokensFilter {
    pub count: usize,
    /// Buffer holding all bytes from the stream.
    buffer: Vec<u8>,
    pub tokenizer: Tokenizer,
    pub encoding: TokenEncoding,
}

impl TailTokensFilter {
    pub fn new(count: usize) -> Self {
        let encoding = TokenEncoding::default();
        Self {
            count,
            buffer: Vec::with_capacity(PIPESIZE),
            tokenizer: get_tokenizer(encoding).clone(),
            encoding,
        }
    }
}

impl FilterPlugin for TailTokensFilter {
    fn filter(&mut self, reader: &mut dyn Read, writer: &mut dyn Write) -> Result<()> {
        if self.count == 0 {
            return Ok(());
        }

        let tokenizer = &self.tokenizer;

        // Buffer all bytes from the stream
        std::io::copy(reader, &mut self.buffer)?;

        if self.buffer.is_empty() {
            return Ok(());
        }

        let text = String::from_utf8_lossy(&self.buffer);
        let token_strs = tokenizer
            .split_by_token(&text)
            .map_err(|e| std::io::Error::other(e.to_string()))?;

        if token_strs.len() <= self.count {
            // All tokens fit — write everything
            writer.write_all(&self.buffer)?;
        } else {
            // Write only the last N tokens
            let skip = token_strs.len() - self.count;
            let mut byte_offset = 0usize;
            for token_str in token_strs.iter().take(skip) {
                byte_offset += token_str.len();
            }
            let write_len = map_lossy_pos_to_bytes(&self.buffer, &text, byte_offset);
            if write_len < self.buffer.len() {
                writer.write_all(&self.buffer[write_len..])?;
            }
        }

        Ok(())
    }

    fn clone_box(&self) -> Box<dyn FilterPlugin> {
        Box::new(Self {
            count: self.count,
            buffer: Vec::new(),
            tokenizer: get_tokenizer(self.encoding).clone(),
            encoding: self.encoding,
        })
    }

    fn options(&self) -> Vec<FilterOption> {
        vec![
            FilterOption {
                name: "count".to_string(),
                default: None,
                required: true,
            },
            FilterOption {
                name: "encoding".to_string(),
                default: Some(serde_json::Value::String("cl100k_base".to_string())),
                required: false,
            },
        ]
    }

    fn description(&self) -> &str {
        "Read the last N LLM tokens"
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Map a byte position in a lossy string back to a position in the original byte slice.
///
/// `String::from_utf8_lossy` replaces invalid UTF-8 bytes with the Unicode
/// replacement character (U+FFFD), which encodes to 3 bytes in UTF-8. This
/// function walks both the original bytes and the lossy string in lockstep,
/// finding the original byte position that corresponds to `lossy_pos`.
fn map_lossy_pos_to_bytes(original: &[u8], lossy: &str, lossy_pos: usize) -> usize {
    if lossy_pos == 0 {
        return 0;
    }

    let replacement = '\u{FFFD}';
    let replacement_len = replacement.len_utf8(); // 3 bytes

    let mut orig_idx = 0usize;
    let mut lossy_idx = 0usize;
    let lossy_bytes = lossy.as_bytes();

    while lossy_idx < lossy_pos && orig_idx < original.len() {
        // Try to decode the next character from the original bytes
        match std::str::from_utf8(&original[orig_idx..]) {
            Ok("") => break,
            Ok(s) => {
                let ch = s.chars().next().unwrap();
                let ch_len = ch.len_utf8();
                // Check if this is a replacement character in the lossy string
                if ch == replacement
                    && lossy_idx + replacement_len <= lossy_pos
                    && lossy_bytes[lossy_idx..].starts_with(
                        &replacement.encode_utf8(&mut [0; 4]).as_bytes()[..replacement_len],
                    )
                {
                    // Could be a real U+FFFD or a replacement of invalid bytes.
                    // If the original byte at this position is valid UTF-8 start, it's real.
                    if original[orig_idx] < 0x80 || original[orig_idx] >= 0xC0 {
                        // Real character
                        orig_idx += ch_len;
                        lossy_idx += ch_len;
                    } else {
                        // Invalid byte that was replaced — advance original by 1
                        orig_idx += 1;
                        lossy_idx += replacement_len;
                    }
                } else {
                    orig_idx += ch_len;
                    lossy_idx += ch_len;
                }
            }
            Err(e) => {
                let valid = e.valid_up_to();
                if valid > 0 {
                    // Some valid bytes, then invalid
                    orig_idx += valid;
                    lossy_idx += valid;
                } else {
                    // Invalid byte — in lossy it becomes 3-byte replacement char
                    orig_idx += 1;
                    lossy_idx += replacement_len;
                }
            }
        }
    }

    orig_idx.min(original.len())
}

// ---------------------------------------------------------------------------
// Registration
// ---------------------------------------------------------------------------

#[ctor::ctor]
fn register_token_filters() {
    register_filter_plugin("head_tokens", || Box::new(HeadTokensFilter::new(0)))
        .expect("Failed to register head_tokens filter");
    register_filter_plugin("skip_tokens", || Box::new(SkipTokensFilter::new(0)))
        .expect("Failed to register skip_tokens filter");
    register_filter_plugin("tail_tokens", || Box::new(TailTokensFilter::new(0)))
        .expect("Failed to register tail_tokens filter");
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    fn make_tokenizer() -> Tokenizer {
        get_tokenizer(TokenEncoding::Cl100kBase).clone()
    }

    #[test]
    fn test_head_tokens_basic() {
        let mut filter = HeadTokensFilter::new(3);
        filter.tokenizer = make_tokenizer();

        let input = b"The quick brown fox";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();

        let result = String::from_utf8_lossy(&output);
        // "The quick brown" is typically 3 tokens
        assert!(!result.is_empty());
        assert!(result.len() <= input.len());
    }

    #[test]
    fn test_head_tokens_zero() {
        let mut filter = HeadTokensFilter::new(0);
        filter.tokenizer = make_tokenizer();

        let input = b"The quick brown fox";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();
        assert!(output.is_empty());
    }

    #[test]
    fn test_head_tokens_more_than_available() {
        let mut filter = HeadTokensFilter::new(1000);
        filter.tokenizer = make_tokenizer();

        let input = b"Hello world";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();
        assert_eq!(output, input);
    }

    #[test]
    fn test_skip_tokens_basic() {
        let mut filter = SkipTokensFilter::new(2);
        filter.tokenizer = make_tokenizer();

        let input = b"The quick brown fox";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();

        let result = String::from_utf8_lossy(&output);
        // Should have skipped some tokens
        assert!(result.len() < input.len());
    }

    #[test]
    fn test_skip_tokens_zero() {
        let mut filter = SkipTokensFilter::new(0);
        filter.tokenizer = make_tokenizer();

        let input = b"Hello world";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();
        assert_eq!(output, input);
    }

    #[test]
    fn test_tail_tokens_basic() {
        let mut filter = TailTokensFilter::new(2);
        filter.tokenizer = make_tokenizer();

        let input = b"The quick brown fox jumps over the lazy dog";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();

        let result = String::from_utf8_lossy(&output);
        // Should only have last 2 tokens
        assert!(!result.is_empty());
        assert!(result.len() < input.len());
    }

    #[test]
    fn test_tail_tokens_zero() {
        let mut filter = TailTokensFilter::new(0);
        filter.tokenizer = make_tokenizer();

        let input = b"Hello world";
        let mut output = Vec::new();
        filter.filter(&mut Cursor::new(input), &mut output).unwrap();
        assert!(output.is_empty());
    }

    #[test]
    fn test_map_lossy_pos_ascii() {
        let original = b"Hello world";
        let lossy = String::from_utf8_lossy(original);
        assert_eq!(map_lossy_pos_to_bytes(original, &lossy, 5), 5);
        assert_eq!(map_lossy_pos_to_bytes(original, &lossy, 0), 0);
        assert_eq!(map_lossy_pos_to_bytes(original, &lossy, 11), 11);
    }

    #[test]
    fn test_map_lossy_pos_with_invalid_utf8() {
        let original = b"Hello\x80world";
        let lossy = String::from_utf8_lossy(original);
        // lossy = "Hello\u{FFFD}world" (13 bytes)
        // Position 5 in lossy = after "Hello" = position 5 in original
        assert_eq!(map_lossy_pos_to_bytes(original, &lossy, 5), 5);
        // Position 8 in lossy = "Hello\u{FFFD}" = position 6 in original
        // (the invalid byte \x80 at position 5 was replaced)
        assert_eq!(map_lossy_pos_to_bytes(original, &lossy, 8), 6);
    }
}
