use crate::{Error, Result};

const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
const MAX_LINE: usize = 8 * 1024;
const MAX_HEX_DIGITS: usize = 8;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ChunkedDecoder {
    buf: Vec<u8>,
    state: ChunkState,
    finished: bool,
    poison: Option<String>,
    trailer_len: usize,
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum ChunkState {
    Size,
    Data(usize),
    DataLineEnd,
    Trailer,
}

impl Default for ChunkedDecoder {
    fn default() -> Self {
        Self::new()
    }
}

impl ChunkedDecoder {
    pub fn new() -> Self {
        Self {
            buf: Vec::new(),
            state: ChunkState::Size,
            finished: false,
            poison: None,
            trailer_len: 0,
        }
    }

    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<u8>> {
        if let Some(message) = &self.poison {
            return Err(Error::Parse(message.clone()));
        }
        if self.finished {
            return Ok(Vec::new());
        }
        self.buf.extend_from_slice(bytes);
        let mut out = Vec::new();

        loop {
            match self.state {
                ChunkState::Size => {
                    let Some(line) = self.read_line()? else {
                        break;
                    };
                    let size = match parse_size_line(&line) {
                        Ok(size) => size,
                        Err(err) => return self.poison_from(err),
                    };
                    if size == 0 {
                        self.state = ChunkState::Trailer;
                    } else {
                        self.state = ChunkState::Data(size);
                    }
                }
                ChunkState::Data(remaining) => {
                    if self.buf.is_empty() {
                        break;
                    }
                    let n = remaining.min(self.buf.len());
                    out.extend_from_slice(&self.buf[..n]);
                    self.buf.drain(..n);
                    let left = remaining - n;
                    self.state = if left == 0 {
                        ChunkState::DataLineEnd
                    } else {
                        ChunkState::Data(left)
                    };
                }
                ChunkState::DataLineEnd => {
                    if self.buf.is_empty() {
                        break;
                    }
                    if self.buf[0] == b'\n' {
                        self.buf.drain(..1);
                        self.state = ChunkState::Size;
                    } else if self.buf[0] == b'\r' {
                        if self.buf.len() < 2 {
                            break;
                        }
                        if self.buf[1] != b'\n' {
                            return self.poison("malformed chunk data terminator");
                        }
                        self.buf.drain(..2);
                        self.state = ChunkState::Size;
                    } else {
                        return self.poison("malformed chunk data terminator");
                    }
                }
                ChunkState::Trailer => {
                    let Some(line) = self.read_line()? else {
                        break;
                    };
                    self.trailer_len += line.len();
                    if self.trailer_len > MAX_LINE {
                        return self.poison("chunk trailer section too long");
                    }
                    if line.is_empty() {
                        self.finished = true;
                        break;
                    }
                }
            }
        }

        Ok(out)
    }

    pub fn finished(&self) -> bool {
        self.finished
    }

    pub fn reset(&mut self) {
        *self = Self::new();
    }

    fn read_line(&mut self) -> Result<Option<Vec<u8>>> {
        if let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
            let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
            if line.ends_with(b"\n") {
                line.pop();
            }
            if line.ends_with(b"\r") {
                line.pop();
            }
            return Ok(Some(line));
        }
        if self.buf.len() > MAX_LINE {
            return self.poison("chunk line too long");
        }
        Ok(None)
    }

    fn poison<T>(&mut self, message: &str) -> Result<T> {
        self.poison = Some(message.to_string());
        Err(Error::Parse(message.to_string()))
    }

    fn poison_from<T>(&mut self, err: Error) -> Result<T> {
        let message = err.to_string();
        self.poison = Some(message.clone());
        Err(Error::Parse(message))
    }
}

fn parse_size_line(line: &[u8]) -> Result<usize> {
    if line.len() > MAX_LINE {
        return Err(Error::Parse("chunk size line too long".into()));
    }
    let size_part = line.split(|&b| b == b';').next().unwrap_or(line);
    let size_part = trim_ascii(size_part);
    if size_part.is_empty() {
        return Err(Error::Parse("empty chunk size".into()));
    }
    if size_part.len() > MAX_HEX_DIGITS {
        return Err(Error::Parse("chunk size has too many hex digits".into()));
    }
    let mut size = 0usize;
    for &b in size_part {
        let digit = match b {
            b'0'..=b'9' => usize::from(b - b'0'),
            b'a'..=b'f' => usize::from(b - b'a' + 10),
            b'A'..=b'F' => usize::from(b - b'A' + 10),
            _ => return Err(Error::Parse("invalid chunk size digit".into())),
        };
        size = (size << 4) | digit;
    }
    if size > MAX_CHUNK_SIZE {
        return Err(Error::Parse("chunk size exceeds limit".into()));
    }
    Ok(size)
}

fn trim_ascii(bytes: &[u8]) -> &[u8] {
    let mut start = 0;
    let mut end = bytes.len();
    while start < end && bytes[start].is_ascii_whitespace() {
        start += 1;
    }
    while end > start && bytes[end - 1].is_ascii_whitespace() {
        end -= 1;
    }
    &bytes[start..end]
}
