uri-template-core 0.1.0

use alloc::{borrow::Cow, string::String, vec::Vec};
use core::ops::Range;

use crate::__private::{Literal, Modifier, Operator, Piece, Variable};
use crate::Values;
use crate::utils::{decode_digit, decode_pct_hex, pct_encode_byte, truncate_str};

#[derive(Debug)]
pub struct UriTemplate {
    pub(crate) raw: Cow<'static, str>,
    pub(crate) pieces: Cow<'static, [Piece]>,
}

#[cfg(test)]
impl UriTemplate {
    pub fn variables(&self) -> Vec<&str> {
        let mut names = Vec::new();
        for piece in self.pieces.iter() {
            if let Piece::Expression { variables, .. } = piece {
                for var in variables.iter() {
                    names.push(&self.raw[var.range.clone()]);
                }
            }
        }
        names
    }

    pub fn literals(&self) -> Vec<&str> {
        let mut literals = Vec::new();
        for piece in self.pieces.iter() {
            if let Piece::Literal(literal) = piece {
                literals.push(self.raw_slice(literal.range()));
            }
        }
        literals
    }
}

impl UriTemplate {
    pub fn new(input: impl Into<Cow<'static, str>>) -> Result<Self, Error> {
        let raw = input.into();
        let mut pieces = Cow::<[Piece]>::Owned(Vec::new());

        let mut literal_start = 0usize;

        let mut chars = raw.char_indices().peekable();
        while let Some((byte_idx, ref ch)) = chars.next() {
            match ch {
                '{' => {
                    // avoid empty strings at the junction of two non-`Literal::Fine` pieces.
                    // such as:
                    // `%7B{var}`
                    // `你好{var}`
                    // `{var1}{var2}`
                    if literal_start < byte_idx {
                        pieces
                            .to_mut()
                            .push(Piece::Literal(Literal::Fine(literal_start..byte_idx)));
                    }
                    // parse expression
                    // parse operator
                    let Some((byte_idx, possible_operator)) = chars.peek() else {
                        return Err(Error::new_parse_error(
                            ParseError::UnmatchedBrace,
                            raw,
                            byte_idx,
                        ));
                    };
                    let mut out_byte_idx = *byte_idx;
                    let operator = match Operator::try_from_char(possible_operator) {
                        Some(Operator::Reserved(_)) => {
                            return Err(Error::new_parse_error(
                                ParseError::ReservedOperator,
                                raw,
                                out_byte_idx,
                            ));
                        }
                        op => op,
                    };

                    let mut var_name_start: usize = if operator.is_some() {
                        // has the operator, skip to next char
                        chars.next();
                        out_byte_idx + 1
                    } else {
                        // without the operator, continue with the current index
                        out_byte_idx
                    };
                    // parse variables
                    let mut variables = Cow::<[Variable]>::Owned(Vec::new());
                    let mut var_name_end = None;
                    let mut modifier: Option<Modifier> = None;
                    while let Some((inner_byte_idx, ref ch)) = chars.next() {
                        out_byte_idx = inner_byte_idx;
                        match ch {
                            '{' => {
                                return Err(Error::new_parse_error(
                                    ParseError::NestedBrace,
                                    raw,
                                    inner_byte_idx,
                                ));
                            }
                            '}' => {
                                let var_name_end = var_name_end.unwrap_or(inner_byte_idx);
                                if var_name_start < var_name_end {
                                    // example: `{*}`
                                    variables.to_mut().push(Variable {
                                        range: var_name_start..(var_name_end),
                                        modifier: modifier.take(),
                                    });
                                } else {
                                    // example: `{}`
                                    return Err(Error::new_parse_error(
                                        ParseError::EmptyVariableName,
                                        raw,
                                        inner_byte_idx,
                                    ));
                                }
                                break;
                            }
                            ',' => {
                                if var_name_start < inner_byte_idx {
                                    variables.to_mut().push(Variable {
                                        range: var_name_start
                                            ..(var_name_end.unwrap_or(inner_byte_idx)),
                                        modifier: modifier.take(),
                                    });
                                    // reset for next variable
                                    var_name_start = inner_byte_idx + 1;
                                    var_name_end = None;
                                    modifier = None;
                                } else {
                                    // example: `{a,}`
                                    return Err(Error::new_parse_error(
                                        ParseError::EmptyVariableName,
                                        raw,
                                        inner_byte_idx,
                                    ));
                                }
                            }
                            '*' => {
                                // modifier: explode
                                var_name_end = Some(inner_byte_idx);
                                match modifier {
                                    None => modifier = Some(Modifier::Explode),
                                    _ => {
                                        // example: `{var:3*}`
                                        return Err(Error::new_parse_error(
                                            ParseError::MultipleModifierInVariable,
                                            raw,
                                            inner_byte_idx,
                                        ));
                                    }
                                }
                            }
                            ':' => {
                                // modifier: prefix
                                var_name_end = Some(inner_byte_idx);
                                let mut acc = 0u16;

                                let mut digits_num = 0;
                                while let Some((byte_idx, ch)) = chars.peek()
                                    && let Ok(digit) = decode_digit(ch)
                                {
                                    digits_num += 1;
                                    if digits_num == 0 && digit == 0 {
                                        let byte_idx = *byte_idx;
                                        return Err(Error::new_parse_error(
                                            ParseError::PrefixLengthLeadingZero,
                                            raw,
                                            byte_idx,
                                        ));
                                    }

                                    // RFC 6570
                                    // max-length < 10000
                                    if digits_num > 4 {
                                        let byte_idx = *byte_idx;
                                        return Err(Error::new_parse_error(
                                            ParseError::PrefixLengthOverflow,
                                            raw,
                                            byte_idx,
                                        ));
                                    }
                                    acc = acc * 10 + (digit as u16);
                                    chars.next();
                                }
                                if acc == 0 {
                                    return Err(Error::new_parse_error(
                                        ParseError::PrefixLengthEmpty,
                                        raw,
                                        inner_byte_idx,
                                    ));
                                }

                                match modifier {
                                    None => modifier = Some(Modifier::Prefix(acc)),
                                    _ => {
                                        return Err(Error::new_parse_error(
                                            ParseError::MultipleModifierInVariable,
                                            raw,
                                            inner_byte_idx,
                                        ));
                                    }
                                }
                            }
                            ch => {
                                // rfc: varname = varchar *( ["."] varchar )
                                // TODO: prohibit leading dots and consecutive dots in variable names
                                if is_varchar(ch) || ch == &'.' {
                                    continue;
                                } else {
                                    return Err(Error::new_parse_error(
                                        ParseError::NotAllowedCharInExpression,
                                        raw,
                                        inner_byte_idx,
                                    ));
                                }
                            }
                        }
                    }
                    if variables.is_empty() {
                        return Err(Error::new_parse_error(
                            ParseError::EmptyVariableName,
                            raw,
                            out_byte_idx,
                        ));
                    }
                    pieces.to_mut().push(Piece::Expression {
                        operator,
                        variables,
                    });
                    // reset literal buffer start
                    literal_start = out_byte_idx + 1;
                }
                '}' => {
                    return Err(Error::new_parse_error(
                        ParseError::UnmatchedBrace,
                        raw,
                        byte_idx,
                    ));
                }
                '%' => {
                    // avoid push redundant empty literals when using templates that start with '%'.
                    // for example: `%7Bvar%7D`
                    if literal_start < byte_idx {
                        pieces
                            .to_mut()
                            .push(Piece::Literal(Literal::Fine(literal_start..byte_idx)));
                    }
                    let mut pct_encoded_buf = Vec::new();
                    loop {
                        if let Some((_, c0)) = chars.next()
                            && let Some((__byte_idx, c1)) = chars.next()
                            && let Ok(byte) = decode_pct_hex(&c0, &c1)
                        {
                            pct_encoded_buf.push(byte);

                            match chars.peek() {
                                Some((_, '%')) => {
                                    // continue to process the next percentage encoding
                                    chars.next();
                                    continue;
                                }
                                _ => {
                                    // complete the percentage encoding processing
                                    pieces.to_mut().push(Piece::Literal(Literal::Encoded {
                                        range: byte_idx..__byte_idx + 1,
                                        decoded: match String::from_utf8(pct_encoded_buf) {
                                            Err(err) => Err(Cow::Owned(err.into_bytes())),
                                            Ok(s) => Ok(Cow::Owned(s)),
                                        },
                                    }));
                                    literal_start = __byte_idx + 1;
                                    break;
                                }
                            }
                        } else {
                            return Err(Error::new_parse_error(
                                ParseError::NotAllowedSinglePercentInLiteral,
                                raw,
                                byte_idx,
                            ));
                        }
                    }
                }
                ch => match LiteralCharKind::from(ch) {
                    LiteralCharKind::Unreserved | LiteralCharKind::Reserved => {}
                    LiteralCharKind::NeedEncode => {
                        if literal_start < byte_idx {
                            pieces
                                .to_mut()
                                .push(Piece::Literal(Literal::Fine(literal_start..byte_idx)));
                        }
                        let mut literal_end = byte_idx + ch.len_utf8();
                        while let Some((byte_idx, ch)) = chars.peek()
                            && LiteralCharKind::from(ch) == LiteralCharKind::NeedEncode
                        {
                            literal_end = *byte_idx + ch.len_utf8();
                            chars.next();
                        }
                        pieces.to_mut().push(Piece::Literal(Literal::NeedEncode {
                            range: byte_idx..literal_end,
                            encoded: {
                                let to_encode = &raw[byte_idx..literal_end];
                                let mut encoded = String::with_capacity(to_encode.len() * 3);
                                for b in to_encode.as_bytes().iter() {
                                    let [c0, c1, c2] = pct_encode_byte(b);
                                    encoded.push(c0);
                                    encoded.push(c1);
                                    encoded.push(c2);
                                }
                                Cow::Owned(encoded)
                            },
                        }));
                        literal_start = literal_end;
                    }
                    LiteralCharKind::Invalid => {
                        return Err(Error::new_parse_error(
                            ParseError::NotAllowedCharInLiteral,
                            raw,
                            byte_idx,
                        ));
                    }
                },
            }
        }

        if literal_start < raw.len() {
            pieces
                .to_mut()
                .push(Piece::Literal(Literal::Fine(literal_start..raw.len())));
        }
        Ok(Self { raw, pieces })
    }
}

impl UriTemplate {
    #[inline]
    fn raw_slice(&self, range: &Range<usize>) -> &str {
        &self.raw[range.clone()]
    }

    pub fn expand<'v, 's: 'v>(&self, values: impl Values<'v, 's>) -> Result<String, Error> {
        let mut out = String::new();
        for piece in self.pieces.iter() {
            match piece {
                Piece::Literal(literal) => match literal {
                    Literal::Fine(range) | Literal::Encoded { range, .. } => {
                        out.push_str(self.raw_slice(range));
                    }
                    Literal::NeedEncode { encoded, .. } => {
                        // https://www.rfc-editor.org/rfc/rfc6570#section-2.1
                        // ..., if not
                        // allowed, copied to the URI reference as the sequence of pct-encoded
                        // triplets corresponding to that character's encoding in UTF-8
                        out.push_str(encoded);
                    }
                },
                Piece::Expression {
                    operator,
                    variables,
                } => {
                    let operator_mode = Operator::mode(operator);
                    let mut first_in_variables = true;

                    #[inline(always)]
                    fn push_str(out: &mut String, s: &str, should_encode_reserved: bool) {
                        if should_encode_reserved {
                            pct_encode_str_reserved(s, out);
                        } else {
                            pct_encode_str_unreserved(s, out);
                        }
                    }
                    for var in variables.iter() {
                        let var_name = self.raw_slice(var.range());
                        if let Some(value) = values.var(var_name)
                            && value.should_expand()
                        {
                            if first_in_variables {
                                // push leading char only if any of the variables are defined.
                                if let Some(leading) = operator_mode.leading {
                                    out.push(leading);
                                }
                                first_in_variables = false;
                            } else {
                                out.push(operator_mode.separator);
                            }
                            match value.as_ref() {
                                crate::Value::Scalar(s) => {
                                    if operator_mode.should_named {
                                        out.push_str(var_name);
                                        if !s.is_empty()
                                            || operator_mode.should_append_equal_when_empty_string
                                        {
                                            out.push('=');
                                        }
                                    }
                                    let s = if let Some(Modifier::Prefix(prefix)) = var.modifier {
                                        truncate_str(s, prefix as usize)
                                    } else {
                                        s
                                    };
                                    push_str(&mut out, s, operator_mode.should_encode_reserved)
                                }
                                crate::Value::List(list) => match var.modifier {
                                    None | Some(Modifier::Prefix(_)) => {
                                        let mut first_in_list = true;
                                        for s in list.iter() {
                                            if first_in_list {
                                                if operator_mode.should_named {
                                                    out.push_str(var_name);
                                                    out.push('=');
                                                }
                                                first_in_list = false;
                                            } else {
                                                out.push(',');
                                            }
                                            let s = if let Some(Modifier::Prefix(prefix)) =
                                                var.modifier
                                            {
                                                truncate_str(s, prefix as usize)
                                            } else {
                                                s
                                            };
                                            push_str(
                                                &mut out,
                                                s,
                                                operator_mode.should_encode_reserved,
                                            );
                                        }
                                    }
                                    Some(Modifier::Explode) => {
                                        let mut first_in_list = true;
                                        for s in list.iter() {
                                            if first_in_list {
                                                first_in_list = false;
                                            } else {
                                                out.push(operator_mode.separator);
                                            }
                                            if operator_mode.should_named {
                                                out.push_str(var_name);
                                                out.push('=');
                                            }
                                            push_str(
                                                &mut out,
                                                s,
                                                operator_mode.should_encode_reserved,
                                            );
                                        }
                                    }
                                },
                                crate::Value::Keys(keys) => match var.modifier {
                                    Some(Modifier::Prefix(_)) => {
                                        return Err(Error::new_expand_error(
                                            "cannot apply prefix modifier to key-value pairs"
                                                .into(),
                                        ));
                                    }
                                    None => {
                                        let mut first_in_keys = true;
                                        for (key, val) in keys.iter() {
                                            if first_in_keys {
                                                if operator_mode.should_named {
                                                    out.push_str(var_name);
                                                    out.push('=');
                                                }
                                                first_in_keys = false;
                                            } else {
                                                out.push(',');
                                            }
                                            let s = if let Some(Modifier::Prefix(prefix)) =
                                                var.modifier
                                            {
                                                truncate_str(key, prefix as usize)
                                            } else {
                                                key
                                            };
                                            push_str(
                                                &mut out,
                                                s,
                                                operator_mode.should_encode_reserved,
                                            );
                                            out.push(',');
                                            let s = if let Some(Modifier::Prefix(prefix)) =
                                                var.modifier
                                            {
                                                truncate_str(val, prefix as usize)
                                            } else {
                                                val
                                            };
                                            push_str(
                                                &mut out,
                                                s,
                                                operator_mode.should_encode_reserved,
                                            );
                                        }
                                    }
                                    Some(Modifier::Explode) => {
                                        let mut first_in_keys = true;
                                        for (key, val) in keys.iter() {
                                            if first_in_keys {
                                                first_in_keys = false;
                                            } else {
                                                out.push(operator_mode.separator);
                                            }
                                            push_str(
                                                &mut out,
                                                key,
                                                operator_mode.should_encode_reserved,
                                            );
                                            out.push('=');
                                            push_str(
                                                &mut out,
                                                val,
                                                operator_mode.should_encode_reserved,
                                            );
                                        }
                                    }
                                },
                            }
                        }
                    }
                }
            }
        }
        Ok(out)
    }
}

impl Operator {
    const fn try_from_char(ch: &char) -> Option<Self> {
        match ch {
            '+' => Some(Self::Plus),
            '/' => Some(Self::Slash),
            '.' => Some(Self::Dot),
            ';' => Some(Self::Semi),
            '?' => Some(Self::Question),
            '&' => Some(Self::Ampersand),
            '#' => Some(Self::Hash),
            '=' | ',' | '!' | '@' | '|' => Some(Self::Reserved(*ch as u8)),
            _ => None,
        }
    }

    const fn mode(this: &Option<Self>) -> OperatorMode {
        const DEFAULT: OperatorMode = OperatorMode {
            leading: None,
            separator: ',',
            should_encode_reserved: true,
            should_named: false,
            should_append_equal_when_empty_string: false,
        };
        let Some(this) = this else {
            return DEFAULT;
        };
        match this {
            Self::Plus => OperatorMode {
                leading: None,
                separator: ',',
                should_encode_reserved: false,
                should_named: false,
                should_append_equal_when_empty_string: false,
            },
            Self::Hash => OperatorMode {
                separator: ',',
                leading: Some('#'),
                should_encode_reserved: false,
                should_named: false,
                should_append_equal_when_empty_string: false,
            },
            Self::Dot => OperatorMode {
                leading: Some('.'),
                separator: '.',
                should_encode_reserved: true,
                should_named: false,
                should_append_equal_when_empty_string: false,
            },
            Self::Slash => OperatorMode {
                leading: Some('/'),
                separator: '/',
                should_encode_reserved: true,
                should_named: false,
                should_append_equal_when_empty_string: false,
            },
            Self::Semi => OperatorMode {
                leading: Some(';'),
                separator: ';',
                should_encode_reserved: true,
                should_named: true,
                should_append_equal_when_empty_string: false,
            },
            Self::Question => OperatorMode {
                leading: Some('?'),
                separator: '&',
                should_encode_reserved: true,
                should_named: true,
                should_append_equal_when_empty_string: true,
            },
            Self::Ampersand => OperatorMode {
                leading: Some('&'),
                separator: '&',
                should_encode_reserved: true,
                should_named: true,
                should_append_equal_when_empty_string: true,
            },
            Self::Reserved(_) => DEFAULT,
        }
    }
}

pub(crate) struct OperatorMode {
    pub leading: Option<char>,
    pub separator: char,
    pub should_encode_reserved: bool,
    pub should_named: bool,
    pub should_append_equal_when_empty_string: bool,
}

#[derive(Debug)]
pub struct Error {
    kind: ErrorKind,
}

impl Error {
    pub fn is_parse_error(&self) -> bool {
        matches!(self.kind, ErrorKind::Parse { .. })
    }

    pub fn is_expand_error(&self) -> bool {
        matches!(self.kind, ErrorKind::Expand(_))
    }
}

#[derive(Debug)]
enum ErrorKind {
    Expand(Cow<'static, str>),
    Parse {
        kind: ParseError,
        raw: Cow<'static, str>,
        position: usize,
    },
}

impl Error {
    pub(crate) const fn new_expand_error(msg: Cow<'static, str>) -> Self {
        Self {
            kind: ErrorKind::Expand(msg),
        }
    }

    pub(crate) const fn new_parse_error(
        kind: ParseError,
        raw: Cow<'static, str>,
        position: usize,
    ) -> Self {
        Self {
            kind: ErrorKind::Parse {
                kind,
                raw,
                position,
            },
        }
    }
}

#[derive(Debug)]
pub enum ParseError {
    /// `{var1{var2}}`
    NestedBrace,
    /// `{var`
    UnmatchedBrace,
    /// `{}`
    EmptyVariableName,
    ReservedOperator,
    /// `{hello:2*}`
    MultipleModifierInVariable,
    // InvalidPercentEncoded,
    // `%` only allowed in percent-encoding
    NotAllowedSinglePercentInLiteral,
    NotAllowedCharInLiteral,
    NotAllowedCharInExpression,
    /// `{var:}`
    PrefixLengthEmpty,
    /// `{var:03}`
    PrefixLengthLeadingZero,
    /// `{var:10000}`
    PrefixLengthOverflow,
}

impl core::error::Error for Error {}
impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match &self.kind {
            ErrorKind::Expand(msg) => write!(f, "expansion error: {}.", msg),
            ErrorKind::Parse {
                kind,
                raw,
                position,
            } => {
                write!(
                    f,
                    "parse error: {} at byte positions {}.\nraw: {}",
                    kind, position, raw
                )
            }
        }
    }
}

impl core::fmt::Display for ParseError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::NestedBrace => write!(f, "nested brace '{{'"),
            Self::UnmatchedBrace => write!(f, "unmatched `{{...` or `...}}`'"),
            Self::EmptyVariableName => write!(f, "empty variable name `{{}}`"),
            Self::ReservedOperator => write!(f, "not allowed reserved operator in expression",),
            Self::NotAllowedSinglePercentInLiteral => write!(f, "not allowed `%` in literal",),
            Self::NotAllowedCharInLiteral => write!(f, "not allowed char in literal"),
            Self::NotAllowedCharInExpression => write!(f, "not allowed char in expression"),
            Self::MultipleModifierInVariable => write!(f, "multiple modifier in variable"),
            Self::PrefixLengthEmpty => write!(f, "empty prefix length"),
            Self::PrefixLengthLeadingZero => write!(f, "leading zero prefix length"),
            Self::PrefixLengthOverflow => write!(f, "prefix length overflows"),
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
enum LiteralCharKind {
    Unreserved,
    Reserved,
    NeedEncode,
    Invalid,
}

impl LiteralCharKind {
    const fn from(ch: &char) -> Self {
        match ch {
            // ALPHA / DIGIT / "-" / "." / "_" / "~"
            'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '.' | '_' | '~' => {
                LiteralCharKind::Unreserved
            }
            // gen-delims
            ':' | '/' | '?' | '#' | '[' | ']' | '@' => LiteralCharKind::Reserved,
            // sub-delims
            // "'': https://www.rfc-editor.org/errata/eid6937
            '!' | '$' | '&' | '\'' | '(' | ')' | '*' | '+' | ',' | ';' | '=' => {
                LiteralCharKind::Reserved
            }
            // CTL
            '\0'..='\x1F' | '\x7F' => LiteralCharKind::Invalid,
            // other not allowed characters
            // "%": reserved for percent-encoding
            ' ' | '"' | '%' | '<' | '>' | '\\' | '^' | '`' | '{' | '|' | '}' => {
                LiteralCharKind::Invalid
            }
            _ => LiteralCharKind::NeedEncode,
        }
    }
}

/// ALPHA / DIGIT / "_" / pct-encoded
const fn is_varchar(ch: &char) -> bool {
    matches!(ch, 'A'..='Z' | 'a'..='z' | '0'..='9' | '_' | '%')
}

const fn is_unreserved(c: char) -> bool {
    matches!(c, 'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '.' | '_' | '~')
}

const fn is_reserved(c: char) -> bool {
    matches!(
        c,
        ':' | '/'
            | '?'
            | '#'
            | '['
            | ']'
            | '@'
            | '!'
            | '$'
            | '&'
            | '\''
            | '('
            | ')'
            | '*'
            | '+'
            | ','
            | ';'
            | '='
    )
}

fn pct_encode_str_reserved(s: &str, out: &mut String) {
    for ch in s.chars() {
        if is_unreserved(ch) {
            out.push(ch);
        } else {
            for b in ch.encode_utf8(&mut [0; 4]).as_bytes() {
                let [c0, c1, c2] = pct_encode_byte(b);
                out.push(c0);
                out.push(c1);
                out.push(c2);
            }
        }
    }
}

/// used for `+` operator and `#` operator
fn pct_encode_str_unreserved(s: &str, out: &mut String) {
    let mut chars = s.char_indices();
    while let Some((_, ch)) = chars.next() {
        if ch == '%' {
            let mut __chars = chars.clone();
            if let Some((_, c0)) = __chars.next()
                && let Some((_, c1)) = __chars.next()
                && let Ok(_) = decode_pct_hex(&c0, &c1)
            {
                // handle percent-encoded sequences
                out.push('%');
                out.push(c0);
                out.push(c1);
                chars = __chars;
            } else {
                // invalid percent-encoding, encode '%' itself
                out.push('%');
                out.push('2');
                out.push('5');
            }
        } else if is_reserved(ch) || is_unreserved(ch) {
            out.push(ch);
        } else {
            for b in ch.encode_utf8(&mut [0; 4]).as_bytes() {
                let [c0, c1, c2] = pct_encode_byte(b);
                out.push(c0);
                out.push(c1);
                out.push(c2);
            }
        }
    }
}

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

    #[test]
    fn test_pct_encode_str_unreserved() {
        let mut s = String::new();
        pct_encode_str_reserved("!", &mut s);
        assert_eq!(s, "%21");
    }
}