use alloc::{string::String, vec::Vec};
use core::str;

use crate::{
    format::Format,
    helpers::{take_buffer_iter, take_byte, take_byte_iter, take_num, take_num_iter},
    Error, Unpackable,
};

/// Unpacks msgpack nil as `None`, or binary data as `Some(&[u8])`.
pub fn unpack_bytes_option_ref(buf: &[u8]) -> Result<(usize, Option<&[u8]>), Error> {
    if buf.first() == Some(&Format::NIL) {
        Ok((1, None))
    } else {
        unpack_bytes(buf).map(|(n, t)| (n, Some(t)))
    }
}

/// Unpacks msgpack nil as `None`, or binary data as `Some(Vec<u8>)`.
pub fn unpack_bytes_option(buf: &[u8]) -> Result<(usize, Option<Vec<u8>>), Error> {
    unpack_bytes_option_ref(buf).map(|(n, t)| (n, t.map(<[u8]>::to_vec)))
}

/// Iterator-based version: unpacks nil as `None`, binary data as `Some(Vec<u8>)`.
pub fn unpack_bytes_option_iter<I>(bytes: I) -> Result<(usize, Option<Vec<u8>>), Error>
where
    I: IntoIterator<Item = u8>,
{
    let mut bytes = bytes.into_iter();
    let first = bytes.next().ok_or(Error::BufferTooShort)?;
    if first == Format::NIL {
        return Ok((1, None));
    }
    let (n, len) = match first {
        Format::BIN8 => (2, take_byte_iter(bytes.by_ref())? as usize),
        Format::BIN16 => (
            3,
            take_num_iter(bytes.by_ref(), u16::from_be_bytes)? as usize,
        ),
        Format::BIN32 => (
            5,
            take_num_iter(bytes.by_ref(), u32::from_be_bytes)? as usize,
        ),
        _ => return Err(Error::UnexpectedFormatTag),
    };
    let data = take_buffer_iter(bytes, len)?;
    Ok((n + len, Some(data)))
}

/// Unpack bytes on binary format
pub fn unpack_bytes(mut buf: &[u8]) -> Result<(usize, &[u8]), Error> {
    let format = take_byte(&mut buf)?;
    let (n, len) = match format {
        Format::BIN8 => (2, take_byte(&mut buf)? as usize),
        Format::BIN16 => (3, take_num(&mut buf, u16::from_be_bytes)? as usize),
        Format::BIN32 => (5, take_num(&mut buf, u32::from_be_bytes)? as usize),
        _ => return Err(Error::UnexpectedFormatTag),
    };
    if buf.len() < len {
        return Err(Error::BufferTooShort);
    }
    Ok((n + len, &buf[..len]))
}

/// Unpack bytes on binary format using an iterator implementation.
pub fn unpack_bytes_iter<I>(bytes: I) -> Result<(usize, Vec<u8>), Error>
where
    I: IntoIterator<Item = u8>,
{
    let mut bytes = bytes.into_iter();
    let format = take_byte_iter(bytes.by_ref())?;
    let (n, len) = match format {
        Format::BIN8 => (2, take_byte_iter(bytes.by_ref())? as usize),
        Format::BIN16 => (
            3,
            take_num_iter(bytes.by_ref(), u16::from_be_bytes)? as usize,
        ),
        Format::BIN32 => (
            5,
            take_num_iter(bytes.by_ref(), u32::from_be_bytes)? as usize,
        ),
        _ => return Err(Error::UnexpectedFormatTag),
    };
    let data = take_buffer_iter(bytes, len)?;
    Ok((n + len, data))
}

pub fn unpack_str(mut buf: &[u8]) -> Result<(usize, &str), Error> {
    let format = take_byte(&mut buf)?;
    let (n, len) = match format {
        0xa0..=0xbf => (1, format as usize & 0x1f),
        Format::STR8 => (2, take_byte(&mut buf)? as usize),
        Format::STR16 => (3, take_num(&mut buf, u16::from_be_bytes)? as usize),
        Format::STR32 => (5, take_num(&mut buf, u32::from_be_bytes)? as usize),
        _ => return Err(Error::UnexpectedFormatTag),
    };
    if buf.len() < len {
        return Err(Error::BufferTooShort);
    }
    let str = str::from_utf8(&buf[..len]).map_err(|_| Error::InvalidUtf8)?;
    Ok((n + len, str))
}

impl Unpackable for String {
    type Error = Error;

    fn unpack_with_ofs(buf: &[u8]) -> Result<(usize, Self), Self::Error> {
        unpack_str(buf).map(|(n, s)| (n, s.into()))
    }

    fn unpack_iter<I>(bytes: I) -> Result<(usize, Self), Self::Error>
    where
        I: IntoIterator<Item = u8>,
    {
        let mut bytes = bytes.into_iter();
        let format = take_byte_iter(bytes.by_ref())?;
        let (n, len) = match format {
            0xa0..=0xbf => (1, format as usize & 0x1f),
            Format::STR8 => (2, take_byte_iter(bytes.by_ref())? as usize),
            Format::STR16 => (
                3,
                take_num_iter(bytes.by_ref(), u16::from_be_bytes)? as usize,
            ),
            Format::STR32 => (
                5,
                take_num_iter(bytes.by_ref(), u32::from_be_bytes)? as usize,
            ),
            _ => return Err(Error::UnexpectedFormatTag),
        };
        let v: Vec<_> = bytes.take(len).collect();
        if v.len() < len {
            return Err(Error::BufferTooShort);
        }
        let s = String::from_utf8(v).map_err(|_| Error::InvalidUtf8)?;
        Ok((n + len, s))
    }
}
