tds-driver 0.1.1

use std::borrow::Cow;

use crate::{
    error::Error,
    sql_read_bytes::SqlReadBytes,
    tds::collation::{decode_oem, encode_oem, Charset, Collation},
    VarLenType,
};

pub(crate) async fn decode<R>(
    src: &mut R,
    ty: VarLenType,
    len: usize,
    collation: Option<Collation>,
) -> crate::Result<Option<Cow<'static, str>>>
where
    R: SqlReadBytes + Unpin,
{
    use VarLenType::*;

    let data = super::plp::decode(src, len).await?;

    match (data, ty) {
        // Codepages other than UTF
        (Some(buf), BigChar) | (Some(buf), BigVarChar) => {
            // TYPE_INFO always reads one for these types.
            let collation = collation
                .as_ref()
                .ok_or_else(|| Error::Protocol("varchar without a collation".into()))?;
            Ok(Some(decode_varchar(buf, collation)?.into()))
        }
        (Some(buf), _) => Ok(Some(decode_nchar(&buf)?.into())),
        _ => Ok(None),
    }
}

/// Decodes `buf` in the collation's code page. An invalid sequence becomes
/// U+FFFD, as in [`decode_utf16le`]: SQL Server stores bytes its code page
/// can't map (0x82 0x20 in a Japanese_CI_AS column), and failing would lose
/// the whole result set. SqlClient decodes with .NET's
/// `Encoding.GetEncoding(codePage)`, whose best-fit fallback substitutes `?`
/// instead.
pub(crate) fn decode_varchar(buf: Vec<u8>, collation: &Collation) -> crate::Result<String> {
    let encoder = match collation.charset()? {
        Charset::Encoding(encoder) => encoder,
        Charset::Oem(table) => return Ok(decode_oem(table, &buf)),
    };

    match encoder.decode_without_bom_handling(&buf).0 {
        // Borrowed means the bytes are valid UTF-8 as they are, so keep them
        // instead of copying.
        Cow::Borrowed(_) => {
            String::from_utf8(buf).map_err(|_| Error::Encoding("invalid sequence".into()))
        }
        Cow::Owned(s) => Ok(s),
    }
}

/// Decodes an nvarchar, nchar, ntext or xml value. The server pads binary
/// converted to these types to an even length, so an odd one is malformed.
pub(crate) fn decode_nchar(buf: &[u8]) -> crate::Result<String> {
    if !buf.len().is_multiple_of(2) {
        return Err(Error::Protocol(
            format!("UTF-16 value of odd length {}", buf.len()).into(),
        ));
    }

    Ok(decode_utf16le(buf))
}

/// Encodes `s` in the collation's code page; a character it lacks is an error.
pub(crate) fn encode_varchar(s: &str, collation: &Collation) -> crate::Result<Vec<u8>> {
    let (encoding, mut encoder) = match collation.charset()? {
        Charset::Encoding(encoding) => (encoding, encoding.new_encoder()),
        Charset::Oem(table) => return encode_oem(table, s),
    };
    let len = encoder
        .max_buffer_length_from_utf8_without_replacement(s.len())
        .ok_or_else(|| Error::Encoding("string too long to encode".into()))?;
    let mut bytes = Vec::with_capacity(len);
    let (res, _) = encoder.encode_from_utf8_to_vec_without_replacement(s, &mut bytes, true);
    match res {
        encoding_rs::EncoderResult::Unmappable(c) => Err(Error::Encoding(
            format!("{c:?} has no representation in {}", encoding.name()).into(),
        )),
        _ => Ok(bytes),
    }
}

/// Decodes UTF-16LE bytes (even length) into a `String` allocated at its exact
/// UTF-8 size. ASCII-only input is copied byte by byte without decoding.
///
/// An unpaired surrogate becomes U+FFFD. SQL Server stores them (it treats
/// nvarchar as UCS-2) and SqlClient returns them as-is in a UTF-16 .NET
/// string; a Rust `String` cannot hold one, and failing would lose the whole
/// result set.
pub(crate) fn decode_utf16le(buf: &[u8]) -> String {
    let (pairs, _) = buf.as_chunks::<2>();
    let units = pairs.iter().map(|&pair| u16::from_le_bytes(pair));

    // A surrogate half counts 2, so a valid pair counts 4; an unpaired half
    // becomes 3 bytes and may make the String grow once.
    let mut utf8_len = 0usize;
    let mut ascii = true;
    for u in units.clone() {
        utf8_len += match u {
            0..=0x7f => 1,
            0x80..=0x7ff | 0xd800..=0xdfff => 2,
            _ => 3,
        };
        ascii &= u < 0x80;
    }

    if ascii {
        let bytes: Vec<u8> = pairs.iter().map(|&[lo, _]| lo).collect();
        return String::from_utf8(bytes).expect("ASCII is UTF-8");
    }

    let mut s = String::with_capacity(utf8_len);
    s.extend(char::decode_utf16(units).map(|c| c.unwrap_or(char::REPLACEMENT_CHARACTER)));
    s
}

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

    fn utf16le(s: &str) -> Vec<u8> {
        s.encode_utf16().flat_map(u16::to_le_bytes).collect()
    }

    #[test]
    fn decode_utf16le_ascii_and_multibyte() {
        assert_eq!(decode_utf16le(&utf16le("")), "");
        assert_eq!(decode_utf16le(&utf16le("hello")), "hello");
        let s = "åäö €𝄞 ascii";
        let decoded = decode_utf16le(&utf16le(s));
        assert_eq!(decoded, s);
        assert_eq!(decoded.capacity(), s.len());
    }

    #[test]
    fn decode_utf16le_replaces_unpaired_surrogates() {
        // Lone high, lone low, reversed pair, then a valid pair.
        let units = [0xd800, b'a' as u16, 0xdc00, 0xde00, 0xd83d, 0xd83d, 0xde00];
        let buf: Vec<u8> = units.iter().flat_map(|u| u.to_le_bytes()).collect();
        assert_eq!(decode_utf16le(&buf), "\u{FFFD}a\u{FFFD}\u{FFFD}\u{FFFD}😀");
    }

    #[tokio::test]
    async fn decode_nvarchar_with_a_lone_surrogate() {
        use crate::sql_read_bytes::test_utils::IntoSqlReadBytes;
        use bytes::{BufMut, BytesMut};

        let mut buf = BytesMut::new();
        buf.put_u16_le(4);
        buf.put_u16_le(b'x' as u16);
        buf.put_u16_le(0xdbff);
        let s = decode(
            &mut buf.into_sql_read_bytes(),
            VarLenType::NVarchar,
            8000,
            None,
        )
        .await
        .unwrap();
        assert_eq!(s.as_deref(), Some("x\u{FFFD}"));
    }

    #[test]
    fn varchar_in_oem_code_pages() {
        let cp437 = Collation::new(0x0409, 32); // SQL_Latin1_General_CP437_CI_AS
        let cp850 = Collation::new(0x0409, 42); // SQL_Latin1_General_CP850_CI_AS
        assert_eq!(
            decode_varchar(vec![b'a', 0x80, 0xb3, 0xe3, 0xff], &cp437).unwrap(),
            "aÇ│π\u{a0}"
        );
        assert_eq!(
            decode_varchar(vec![b'a', 0x9d, 0xd0, 0xd5, 0xf2], &cp850).unwrap(),
            "aØðı‗"
        );

        assert_eq!(
            encode_varchar("aÇ│π\u{a0}", &cp437).unwrap(),
            [b'a', 0x80, 0xb3, 0xe3, 0xff]
        );
        assert_eq!(
            encode_varchar("aØðı‗", &cp850).unwrap(),
            [b'a', 0x9d, 0xd0, 0xd5, 0xf2]
        );
        assert!(matches!(
            encode_varchar("Ø", &cp437),
            Err(Error::Encoding(_))
        ));

        // Every byte round-trips.
        for c in [cp437, cp850] {
            let all: Vec<u8> = (0..=255).collect();
            let s = decode_varchar(all.clone(), &c).unwrap();
            assert_eq!(s.chars().count(), 256);
            assert_eq!(encode_varchar(&s, &c).unwrap(), all);
        }
    }

    #[test]
    fn varchar_in_a_utf8_collation() {
        // Latin1_General_100_CI_AS_SC_UTF8: LCID 0x409 with fUTF8 set.
        let utf8 = Collation::new(0x0409 | 0x0400_0000, 0);
        assert_eq!(
            decode_varchar("héllo 日😀".as_bytes().to_vec(), &utf8).unwrap(),
            "héllo 日😀"
        );
        assert_eq!(encode_varchar("héllo", &utf8).unwrap(), "héllo".as_bytes());
    }

    #[test]
    fn invalid_code_page_sequences_are_replaced() {
        let sjis = Collation::new(0x0411, 0); // Japanese_CI_AS
        let utf8 = Collation::new(0x0409 | 0x0400_0000, 0);
        // A lead byte before an invalid trail, then one at the end.
        assert_eq!(
            decode_varchar(vec![b'A', 0x82, 0x20, 0x82, 0xa0, 0x82], &sjis).unwrap(),
            "A\u{FFFD} あ\u{FFFD}"
        );
        assert_eq!(
            decode_varchar(vec![b'a', 0xff, 0xe6, 0x97], &utf8).unwrap(),
            "a\u{FFFD}\u{FFFD}"
        );
    }

    #[tokio::test]
    async fn varchar_without_a_collation_is_an_error() {
        use crate::sql_read_bytes::test_utils::IntoSqlReadBytes;
        use bytes::{BufMut, BytesMut};

        let mut buf = BytesMut::new();
        buf.put_u16_le(1);
        buf.put_u8(b'a');
        let res = decode(
            &mut buf.into_sql_read_bytes(),
            VarLenType::BigVarChar,
            10,
            None,
        )
        .await;
        assert!(matches!(res, Err(Error::Protocol(_))), "{res:?}");
    }

    #[test]
    fn odd_length_utf16_is_an_error() {
        assert_eq!(decode_nchar(&[b'h', 0, b'i', 0]).unwrap(), "hi");
        assert!(matches!(
            decode_nchar(&[b'h', 0, b'i']),
            Err(Error::Protocol(_))
        ));
    }

    #[test]
    fn decode_varchar_keeps_ascii_and_converts_the_rest() {
        let collation = Collation::new(0x0409, 0); // Windows-1252
        assert_eq!(
            decode_varchar(b"plain".to_vec(), &collation).unwrap(),
            "plain"
        );
        assert_eq!(
            decode_varchar(vec![0xe5, 0xe4, 0xf6], &collation).unwrap(),
            "åäö"
        );
    }
}