"""
String manipulation functions.

This module provides a comprehensive set of string operations including:
- String conversion and concatenation
- Case conversion and whitespace handling
- String splitting and substring extraction
- Distance metrics
- Pattern matching and searching
- Regular expression matching
"""
from __future__ import annotations

from relationalai.semantics import Float
from relationalai.util.docutils import include_in_docs

from . import DateValue, StringValue, IntegerValue, FloatValue, NumberValue, DateTimeValue
from ..frontend.base import Library, Expression, Field, TupleVariable, Variable
from ..frontend.core import Date, DateTime, Number, String, Integer, Any
from typing import Sequence

__include_in_docs__ = True

# the front-end library object
library = Library("strings")

#--------------------------------------------------
# Operations
#--------------------------------------------------

_string = library.Relation("string", [Field.input("s", Any), Field("result", String)], overloads=[
    [Number, String],
    [Float, String],
    [String, String],
    [DateTime, String],
    [Date, String],
])

@include_in_docs
def string(s: StringValue|IntegerValue|FloatValue|NumberValue|DateTimeValue|DateValue) -> Expression:
    """
    Convert a value to a string representation.

    Parameters
    ----------
    s: StringValue|IntegerValue|FloatValue|NumberValue|DateTimeValue|DateValue
        The value to convert.

    Returns
    -------
    Expression
        An `Expression` representing the `String` conversion of the input. Returns `String`.

    Examples
    --------
    Convert various types to strings:

    >>> strings.string(123)
    >>> strings.string(3.14)
    >>> strings.string(date(2024, 1, 15))
    """
    return _string(s)

_concat = library.Relation("concat", [Field.input("s1", String), Field.input("s2", String), Field("result", String)])

@include_in_docs
def concat(s1: StringValue, s2: StringValue, *sn: StringValue) -> Expression:
    """
    Concatenate multiple strings together.

    Parameters
    ----------
    s1: StringValue
        The first string to concatenate.
    s2: StringValue
        The second string to concatenate.
    *sn: StringValue
        Additional strings to concatenate.

    Returns
    -------
    Expression
        An `Expression` representing the concatenated string. Returns `String`.

    Examples
    --------
    Concatenate strings together:

    >>> strings.concat("thing_", Thing.name)
    >>> strings.concat("Hello", " ", "World")
    """
    res = _concat(s1, s2, String.ref("res0"))
    for i, s in enumerate(sn):
        res = _concat(res, s, String.ref(f"res{i + 1}"))
    return res

_join = library.Relation("join", [Field.input("strs", String, is_list=True), Field.input("sep", String), Field("result", String)])

@include_in_docs
def join(strs: Sequence[StringValue], separator: StringValue = "") -> Expression:
    """
    Join a sequence of strings with a separator.

    Parameters
    ----------
    strs: Sequence[StringValue]
        A sequence of string values to join.
    separator: str
        The separator to use between strings.

    Returns
    -------
    Expression
        An `Expression` representing the joined string. Returns `String`.

    Examples
    --------
    Join strings with various separators:

    >>> strings.join(["FIRST", "MIDDLE", "LAST"]) # defaults to empty string separator
    >>> strings.join([Person.first, Person.last], " ")
    >>> strings.join([Person.first, "the best", Person.last], " | ")
    """
    return _join(TupleVariable(strs), separator)

_contains = library.Relation("contains", [Field.input("s", String), Field.input("substr", String)])

@include_in_docs
def contains(s: StringValue, substr: StringValue) -> Expression:
    """
    Check whether `substr` is contained within `s`.

    Parameters
    ----------
    s: StringValue
        The string to search in.
    substr: StringValue
        The substring to search for.

    Returns
    -------
    Expression
        An `Expression` that evaluates to `true` if `substr` is found in `s`.

    Examples
    --------
    Check if a string contains a substring:

    >>> select(Person.name).where(strings.contains(Person.last_name, "ohn"))
    """
    return _contains(s, substr)

_ends_with = library.Relation("ends_with", [Field.input("s", String), Field.input("suffix", String)])


_starts_with = library.Relation("starts_with", [Field.input("s", String), Field.input("prefix", String)])

@include_in_docs
def startswith(s: StringValue, prefix: StringValue) -> Expression:
    """
    Check whether `s` starts with `prefix`.

    Parameters
    ----------
    s: StringValue
        The string to check.
    prefix: StringValue
        The prefix to check for.

    Returns
    -------
    Expression
        An `Expression` that evaluates to `true` if `s` starts with `prefix`.

    Examples
    --------
    Filter strings that start with a prefix:

    >>> select(Person.name).where(strings.startswith(Person.last_name, Person.name))
    """
    return _starts_with(s, prefix)

@include_in_docs
def endswith(s: StringValue, suffix: StringValue) -> Expression:
    """
    Check whether `s` ends with `suffix`.

    Parameters
    ----------
    s: StringValue
        The string to check.
    suffix: StringValue
        The suffix to check for.

    Returns
    -------
    Expression
        An `Expression` that evaluates to `true` if `s` ends with `suffix`.

    Examples
    --------
    Filter strings that end with a suffix:

    >>> select(Person.name).where(strings.endswith(Person.last_name, "dottir"))
    """
    return _ends_with(s, suffix)

_len = library.Relation("len", [Field.input("s", String), Field("result", Integer)])

@include_in_docs
def len(s: StringValue) -> Expression:
    """
    Get the length of the string `s`.

    Parameters
    ----------
    s: StringValue
        The string to measure.

    Returns
    -------
    Expression
        An `Expression` representing the length of the string. Returns `Integer`.

    Examples
    --------
    Get the length of a string:

    >>> select(Person.name, strings.len(Person.name))
    """
    return _len(s)

_levenshtein = library.Relation("levenshtein", [Field.input("s1", String), Field.input("s2", String), Field("result", Integer)])

@include_in_docs
def levenshtein(s1: StringValue, s2: StringValue) -> Expression:
    """
    Compute the Levenshtein distance between two strings.

    Parameters
    ----------
    s1: StringValue
        The first string.
    s2: StringValue
        The second string.

    Returns
    -------
    Expression
        An `Expression` representing the Levenshtein edit distance. Returns `Integer`.

    Examples
    --------
    Calculate edit distance between two strings:

    >>> strings.levenshtein("kitten", "sitting")
    >>> strings.levenshtein(Car.make, "target")
    """
    return _levenshtein(s1, s2)

_like = library.Relation("like", [Field.input("s", String), Field.input("pattern", String)])

@include_in_docs
def like(s: StringValue, pattern: StringValue) -> Expression:
    """
    Check whether `s` matches the SQL LIKE pattern `pattern`.

    Parameters
    ----------
    s: StringValue
        The string to match against.
    pattern: StringValue
        The SQL LIKE pattern (`%` for wildcard, `_` for single character).

    Returns
    -------
    Expression
        An `Expression` that evaluates to `true` if `s` matches the pattern.

    Examples
    --------
    Match strings using SQL LIKE patterns:

    >>> select(Brand.name).where(strings.like(Brand.name, r"Mc%"))
    >>> select(Brand.name).where(strings.like(Brand.name, r"%est%"))
    """
    return _like(s, pattern)

_lower = library.Relation("lower", [Field.input("s", String), Field("result", String)])

@include_in_docs
def lower(s: StringValue) -> Expression:
    """
    Convert the string `s` to lowercase.

    Parameters
    ----------
    s: StringValue
        The string to convert.

    Returns
    -------
    Expression
        An `Expression` representing the lowercase version of the string. Returns `String`.

    Examples
    --------
    Convert a string to lowercase:

    >>> select(strings.lower(Person.last_name))
    >>> select(City.name).where(strings.lower(City.province) == "ontario")
    """
    return _lower(s)

_upper = library.Relation("upper", [Field.input("s", String), Field("result", String)])

@include_in_docs
def upper(s: StringValue) -> Expression:
    """
    Convert the string `s` to uppercase.

    Parameters
    ----------
    s: StringValue
        The string to convert.

    Returns
    -------
    Expression
        An `Expression` representing the uppercase version of the string. Returns `String`.

    Examples
    --------
    Convert a string to uppercase:

    >>> select(strings.upper(Person.last_name))
    >>> select(City.name).where(strings.upper(City.province) == "ONTARIO")
    """
    return _upper(s)

_replace = library.Relation("replace", [Field.input("source", String), Field.input("old", String), Field.input("new", String), Field("result", String)])

@include_in_docs
def replace(source: StringValue, old: StringValue, new: StringValue) -> Expression:
    """
    Replace occurrences of `old` with `new` in the string `source`.

    Parameters
    ----------
    source: StringValue
        The original string.
    old: StringValue
        The substring to replace.
    new: StringValue
        The replacement substring.

    Returns
    -------
    Expression
        An `Expression` representing the string with replacements applied. Returns `String`.

    Examples
    --------
    Replace substrings within a string:

    >>> strings.replace("My-name", "-", "_")
    >>> strings.replace(Brand.name, "bee", "")
    """
    return _replace(source, old, new)

_split = library.Relation("split", [Field.input("s", String), Field.input("separator", String), Field("index", Integer), Field("result", String)])

@include_in_docs
def split(s: StringValue, separator: StringValue) -> tuple[Variable, Variable]:
    """
    Split the string `s` by `separator`, returning variables holding index and result.

    Parameters
    ----------
    s: StringValue
        The string to split.
    separator: StringValue
        The separator to split by.

    Returns
    -------
    tuple[Variable, Variable]
        A tuple of `(index, result)` variables representing the split parts. Index is `Integer`, result is `String`.

    Examples
    --------
    Split a string and get index and result variables:

    >>> idx, part = strings.split("my_email@mail.test", "@")
    >>> select(idx, part)
    """
    idx = Integer.ref("index")
    res = String.ref("result")
    exp = _split(s, separator, idx, res)
    return exp[2], exp[3]

@include_in_docs
def split_part(s: StringValue, separator: StringValue, index: IntegerValue) -> Expression:
    """
    Get the part of the string `s` at `index` after splitting by `separator`.

    Parameters
    ----------
    s: StringValue
        The string to split.
    separator: StringValue
        The separator to split by.
    index: IntegerValue
        The index of the part to retrieve (0-based).

    Returns
    -------
    Expression
        An `Expression` representing the string part at the specified index. Returns `String`.

    Examples
    --------
    Get a specific part of a split string:

    >>> strings.split_part("my_email@mail.test", "@", 1)
    """
    return _split(s, separator, index)

_strip = library.Relation("strip", [Field.input("s", String), Field("result", String)])

@include_in_docs
def strip(s: StringValue) -> Expression:
    """
    Strip whitespace from both ends of the string `s`.

    Parameters
    ----------
    s: StringValue
        The string to strip.

    Returns
    -------
    Expression
        An `Expression` representing the string with leading and trailing whitespace removed. Returns `String`.

    Examples
    --------
    Remove whitespace from both ends of a string:

    >>> select(strings.strip(Person.name))
    >>> select(Text.content).where(strings.strip(Text.content) == "string with spaces")
    """
    return _strip(s)

_substring = library.Relation("substring", [Field.input("s", String), Field.input("start", Integer), Field.input("stop", Integer), Field("result", String)])

@include_in_docs
def substring(s: StringValue, start: IntegerValue, stop: IntegerValue) -> Expression:
    """
    Get the substring of `s` from `start` to `stop` (0-based, stop exclusive).

    Parameters
    ----------
    s: StringValue
        The string to extract from.
    start: IntegerValue
        The starting index (0-based, inclusive).
    stop: IntegerValue
        The ending index (0-based, exclusive).

    Returns
    -------
    Expression
        An `Expression` representing the extracted substring. Returns `String`.

    Examples
    --------
    Extract a substring from a string:

    >>> strings.substring(Person.name, 0, 3)
    >>> strings.substring("example", 2, 5)
    """
    return _substring(s, start, stop)

_regex_match = library.Relation("regex_match", [Field.input("regex", String), Field.input("value", String)])

@include_in_docs
def regex_match(value: StringValue, regex: StringValue) -> Expression:
    """
    Check if the `value` matches the given `regex`.

    Parameters
    ----------
    value: StringValue
        The string value to match against.
    regex: StringValue
        The regular expression pattern.

    Returns
    -------
    Expression
        An `Expression` that evaluates to `true` if `value` matches the `regex` pattern.

    Examples
    --------
    Match strings using regular expressions:

    >>> select(Person.name).where(strings.regex_match(Person.name, r".*ohn.*"))
    >>> select(Person.name).where(not_(strings.regex_match("ohn", Person.name)))
    """
    return _regex_match(regex, value)
