//! String manipulation utilities.
//!
//! ## Rust rewrite notes
//!
//! These are more or less bug-for-bug reimplementations of the original WUtils
//! functions. Avoid using these functions in new code.

use std::{
    ffi::{c_char, c_int, CStr},
    mem, ptr, slice,
};

use crate::memory::{alloc_bytes, alloc_string, ffi::wrealloc, free_bytes};

/// Returns a `wmalloc`-managed C-style string that duplicates `s`, which cannot
/// be null.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wstrdup(s: *const c_char) -> *mut c_char {
    assert!(!s.is_null());
    alloc_string(unsafe { CStr::from_ptr(s) })
}

/// Returns a `wmalloc`-managed C-style string of the first `len` bytes of `s`,
/// which cannot be null.
///
/// If `len` exceeds the length of `s`, uses the lesser value.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wstrndup(s: *const c_char, len: usize) -> *mut c_char {
    assert!(!s.is_null());
    let len = unsafe {
        slice::from_raw_parts(s, len)
            .into_iter()
            .position(|p| *p == 0)
            .unwrap_or(len)
    };
    let copy: *mut c_char = alloc_bytes(len + 1).cast(); // Implicitly zeroed.
    unsafe {
        ptr::copy_nonoverlapping(s, copy, len);
    }
    copy.cast::<c_char>()
}

/// Concatenates `s1` and `s2` into a `wmalloc`-managed C-style string.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wstrconcat(s1: *const c_char, s2: *const c_char) -> *mut c_char {
    match (s1.is_null(), s2.is_null()) {
        (true, true) => ptr::null_mut(),
        (true, false) => unsafe { wstrdup(s1) },
        (false, true) => unsafe { wstrdup(s2) },
        (false, false) => unsafe {
            let s1 = CStr::from_ptr(s1);
            let l1 = s1.count_bytes();
            let s2 = CStr::from_ptr(s2);
            let l2 = s2.count_bytes();
            let s: *mut c_char = alloc_bytes(l1 + l2 + 1).cast(); // Implicitly zeroed.
            ptr::copy_nonoverlapping(s1.as_ptr(), s, l1);
            ptr::copy_nonoverlapping(s2.as_ptr(), s.offset(l1 as isize), l2);
            s
        },
    }
}

/// Appends `src` to `dest`, destructively reallocating `dest` and returning a
/// new `wmalloc`-managed pointer where the result is stored. `dst` must come
/// from `wmalloc` or `wrealloc`.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wstrappend(dst: *mut c_char, src: *const c_char) -> *mut c_char {
    if src.is_null() {
        return dst;
    }
    let src = unsafe { CStr::from_ptr(src) };
    let src_len = src.count_bytes();
    if src_len == 0 {
        return dst;
    }

    if dst.is_null() {
        return unsafe { wstrdup(src.as_ptr()) };
    }
    let dst_len = unsafe { CStr::from_ptr(dst).count_bytes() + 1 };

    let len = dst_len + src_len + 1;
    let result: *mut c_char = unsafe { wrealloc(dst.cast(), len).cast() };
    unsafe {
        ptr::copy_nonoverlapping(
            src.as_ptr(),
            result.offset(dst_len.try_into().unwrap()),
            src_len,
        );
    }
    result
}

/// Strips leading and trailing whitespace from `s`, returning a
/// `wmalloc`-managed C-style string holding the result.
///
/// ## Rust rewite notes
///
/// This uses a slightly different notion of "space character" than the original
/// C implementation did. `s` must be valid UTF-8.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wtrimspace(s: *const c_char) -> *mut c_char {
    if s.is_null() {
        return ptr::null_mut();
    }

    let Ok(s) = (unsafe { CStr::from_ptr(s).to_str() }) else {
        // TODO: complain.
        return ptr::null_mut();
    };
    let trimmed = s.trim();
    let ptr = trimmed.as_ptr();
    let len = trimmed.len();
    unsafe { wstrndup(ptr, len) }
}

/// Splits `command` into tokens with approximately the same rules as used in
/// the shell for splitting up a command into program arguments.
fn tokensplit(command: &[u8]) -> Vec<Vec<u8>> {
    enum Mode {
        Start,
        Unquoted(Vec<u8>),
        Escape(Context),
        Quoted { token: Vec<u8>, delimiter: u8, },
    }
    enum Context {
        Unquoted(Vec<u8>),
        Quoted { token: Vec<u8>, delimiter: u8, },
    }
    let mut out = Vec::new();
    let mut mode = Mode::Start;
    for &b in command {
        mode = match (mode, b) {
            (Mode::Start, b'\\') => Mode::Escape(Context::Unquoted(vec![])),
            (Mode::Start, b'\'' | b'"') => Mode::Quoted { token: vec![], delimiter: b },
            (Mode::Start, b' ' | b'\t') => Mode::Start,
            (Mode::Start, _) => Mode::Unquoted(vec![b]),
            (Mode::Unquoted(token), b'\\') => Mode::Escape(Context::Unquoted(token)),
            (Mode::Unquoted(token), b'\'' | b'"') => Mode::Quoted { token, delimiter: b },
            (Mode::Unquoted(token), b' ' | b'\t') => {
                out.push(token);
                Mode::Start
            }
            (Mode::Unquoted(mut token), _) => {
                token.push(b);
                Mode::Unquoted(token)
            }
            (Mode::Escape(Context::Unquoted(mut token)), _) => {
                token.push(b);
                Mode::Unquoted(token)
            }
            (Mode::Escape(Context::Quoted { mut token, delimiter }), _) => {
                token.push(b);
                Mode::Quoted { token, delimiter }
            }
            (Mode::Quoted { token, delimiter }, _) if b == delimiter => {
                Mode::Unquoted(token)
            }
            (Mode::Quoted { token, delimiter }, _) if b == b'\\' => {
                Mode::Escape(Context::Quoted { token, delimiter })
            }
            (Mode::Quoted { mut token, delimiter }, _) => {
                token.push(b);
                Mode::Quoted { token, delimiter }
            }
        }
    }
    match mode {
        Mode::Start => (),
        Mode::Unquoted(token) => out.push(token),
        Mode::Quoted { token, .. } => out.push(token),
        Mode::Escape(Context::Unquoted(token)) => out.push(token),
        Mode::Escape(Context::Quoted { token, .. }) => out.push(token),
    }
    out
}

/// Splits `command` into tokens, storing the number of tokens in `argc` and
/// `wmalloc`-managed C-style strings for each token in `argv`. Call
/// [`wtokenfree`] to free `argv`.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wtokensplit(
    command: *const c_char,
    argv: *mut *mut *mut c_char,
    argc: *mut c_int,
) {
    if argv.is_null() || argc.is_null() {
        return;
    }
    unsafe {
        *argv = ptr::null_mut();
        *argc = 0;
    }
    if command.is_null() {
        return;
    }
    let command = unsafe { CStr::from_ptr(command) };
    let Ok(command) = command.to_str() else {
        return;
    };
    let tokens = tokensplit(command.as_bytes());
    if tokens.is_empty() {
        return;
    }

    let argv = unsafe {
        *argv = alloc_bytes(mem::size_of::<*mut c_char>() * tokens.len()).cast::<*mut c_char>();
        *argc = tokens.len() as c_int;
        slice::from_raw_parts_mut(*argv, tokens.len())
    };
    for (dest, mut token) in argv.iter_mut().zip(tokens.into_iter()) {
        token.push(b'\0');
        *dest = alloc_string(unsafe { CStr::from_bytes_with_nul_unchecked(&token) });
    }
}

/// Frees an `argv` populated by [`wtokensplit`].
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wtokenfree(tokens: *mut *mut c_char, count: c_int) {
    if tokens.is_null() {
        return;
    }
    if count > 0 {
        let tokens = unsafe { slice::from_raw_parts_mut(tokens, count as usize) };
        for token in tokens {
            unsafe {
                free_bytes(*token);
            }
        }
    }
    unsafe {
        free_bytes(tokens.cast::<u8>());
    }
}

/// Joins the tokens of the `count` elements of `list` into a single string,
/// enclosing them in double quotes (`"`) if necessary. Returns a
/// `wmalloc`-managed C-style string with the result, or null on failure.
///
/// ## Rust rewrite notes
///
/// `list` must consist of valid UTF-8 strings. This reimplementation does not
/// attempt to improve on the original, so it has bad failure modes (like not
/// escaping quotes in the input). Do not use this function in new code.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wtokenjoin(list: *const *const char, count: c_int) -> *mut c_char {
    if list.is_null() || count <= 0 {
        return alloc_string(c"");
    }

    let list = unsafe {
        slice::from_raw_parts(list.cast::<*const u8>(), count as usize)
    };

    let mut buffer = Vec::new();
    for term in list {
        if term.is_null() {
            continue;
        }
        let term = unsafe { CStr::from_ptr(*term) };
        if term.is_empty() {
            continue;
        }
        if !buffer.is_empty() {
            buffer.push(b' ');
        }
        let term = term.to_bytes();
        if term.iter().find(|&&x| x == b' ' || x == b'\t').is_some() {
            buffer.push(b'"');
            buffer.extend_from_slice(term);
            buffer.push(b'"');
        } else {
            buffer.extend_from_slice(term);
        }
    }
    buffer.push(b'\0');

    if let Ok(buffer) = CStr::from_bytes_until_nul(&buffer) {
        alloc_string(buffer)
    } else {
        return ptr::null_mut();
    }
}

#[cfg(test)]
mod test {
    use super::{wtokenfree, wtokensplit, wtokenjoin};

    use std::{ffi::{c_char, c_int, CStr}, ptr, slice};

    #[test]
    fn split_empty_whitespace() {
        let mut argv = ptr::null_mut();
        let mut argc = 0;

        unsafe { wtokensplit(c"".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 0);
        assert_eq!(argv, ptr::null_mut());
        unsafe { wtokenfree(argv, argc); }

        unsafe { wtokensplit(c"  ".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 0);
        assert_eq!(argv, ptr::null_mut());
        unsafe { wtokenfree(argv, argc); }

        unsafe { wtokensplit(c" \t ".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 0);
        assert_eq!(argv, ptr::null_mut());
        unsafe { wtokenfree(argv, argc); }

        unsafe { wtokensplit(c"  \t ".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 0);
        assert_eq!(argv, ptr::null_mut());
        unsafe { wtokenfree(argv, argc); }

        unsafe { wtokensplit(c"\t\t".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 0);
        assert_eq!(argv, ptr::null_mut());
        unsafe { wtokenfree(argv, argc); }
    }

    fn args_of(argv: *mut *mut u8, argc: usize) -> Vec<String> {
        let mut v = Vec::with_capacity(argc);
        for s in unsafe { slice::from_raw_parts(argv, argc) } {
            if s.is_null() {
                return v;
            }
            v.push(String::from(unsafe { CStr::from_ptr(*s) }.to_str().unwrap()));
        }
        v
    }

    #[test]
    fn split_empty_quoted() {
        let mut argv = ptr::null_mut();
        let mut argc = 0;

        unsafe { wtokensplit(c"\"\"".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("")]);
        unsafe { wtokenfree(argv, argc); }

        argv = ptr::null_mut();
        argc = 0;

        unsafe { wtokensplit(c"''".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("")]);
        unsafe { wtokenfree(argv, argc); }

        argv = ptr::null_mut();
        argc = 0;

        unsafe { wtokensplit(c"  ''\t".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("")]);
        unsafe { wtokenfree(argv, argc); }
    }

    #[test]
    fn split_one_unquoted() {
        let mut argv = ptr::null_mut();
        let mut argc = 0;

        unsafe { wtokensplit(c"hello".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("hello")]);

        unsafe { wtokensplit(c"hello\\\\".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("hello\\")]);
    }

    #[test]
    fn split_one_quoted() {
        let mut argv = ptr::null_mut();
        let mut argc = 0;

        unsafe { wtokensplit(c"\"hello\"".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("hello")]);

        unsafe { wtokensplit(c"\"hello world\"".as_ptr(), &mut argv, &mut argc); }
        assert_eq!(argc, 1);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("hello world")]);
    }

    #[test]
    fn split_multi() {
        let mut argv = ptr::null_mut();
        let mut argc = 0;

        unsafe {
            wtokensplit(
                c"\"hello world\" what\\'s happening' here it\\'s weird'".as_ptr(),
                &mut argv,
                &mut argc,
            );
        }
        assert_eq!(argc, 3);
        assert_eq!(args_of(argv, argc as usize), vec![String::from("hello world"),
                                                      String::from("what's"),
                                                      String::from("happening here it's weird")]);
    }

    /// Calls `f(wtokenjoin(list))` transparently.
    fn with_list<R>(list: &[&CStr], f: impl FnOnce(Option<&CStr>) -> R) -> R {
        let list: Vec<_> = list.iter().map(|s| s.as_ptr().cast::<c_char>()).collect(); 
        let joined = unsafe { wtokenjoin(list.as_slice().as_ptr().cast(), list.len() as c_int) };
        let result = f(if joined.is_null() { None } else { unsafe { Some(CStr::from_ptr(joined)) } });
        unsafe { crate::memory::free_bytes(joined.cast()); }
        result
    }

    #[test]
    fn join_nothing() {
        with_list(&[], |joined| {
            assert_eq!(joined.unwrap(), c"");
        });

        with_list(&[c"", c""], |joined| {
            assert_eq!(joined.unwrap(), c"");
        });
    }

    #[test]
    fn join_basic() {
        with_list(&[c"hello", c"world"], |joined| {
            assert_eq!(joined.unwrap(), c"hello world");
        });
    }

    #[test]
    fn join_quoted() {
        with_list(&[c"hello", c"there world"], |joined| {
            assert_eq!(joined.unwrap(), c"hello \"there world\"");
        });
    }
}
