from __future__ import annotations

from typing import TYPE_CHECKING, Any

from narwhals._compliant.any_namespace import StringNamespace
from narwhals._pandas_like.utils import (
    PandasLikeSeriesNamespace,
    align_and_extract_native,
    get_dtype_backend,
    is_dtype_pyarrow,
)

if TYPE_CHECKING:
    from narwhals._pandas_like.series import PandasLikeSeries


class PandasLikeSeriesStringNamespace(
    PandasLikeSeriesNamespace, StringNamespace["PandasLikeSeries"]
):
    def _with_native_bool(self, result: Any) -> PandasLikeSeries:
        # See https://narwhals-dev.github.io/narwhals/concepts/boolean/.
        if (
            result.dtype == object
            and get_dtype_backend(self.native.dtype, self.implementation) is None
        ):
            not_na = result.notna()
            all_valid = not_na.all()
            if not all_valid and self.implementation._backend_version() < (3,):
                result.where(not_na, False, inplace=True)
                result = result.astype(bool, copy=False)
            elif not all_valid:
                result = result.where(not_na, False).astype(bool)
            else:  # pragma: no cover
                pass
        return self.with_native(result)

    def len_chars(self) -> PandasLikeSeries:
        return self.with_native(self.native.str.len())

    def replace(
        self, value: PandasLikeSeries, pattern: str, *, literal: bool, n: int
    ) -> PandasLikeSeries:
        _, value_native = align_and_extract_native(self.compliant, value)
        if not isinstance(value_native, str):
            msg = f"{self.compliant._implementation} backed `.str.replace` only supports str replacement values"
            raise TypeError(msg)
        series = self.native.str.replace(
            pat=pattern, repl=value_native, n=n, regex=not literal
        )
        return self.with_native(series)

    def replace_all(
        self, value: PandasLikeSeries, pattern: str, *, literal: bool
    ) -> PandasLikeSeries:
        return self.replace(value, pattern, literal=literal, n=-1)

    def strip_chars(self, characters: str | None) -> PandasLikeSeries:
        return self.with_native(self.native.str.strip(characters))

    def strip_chars_start(self, characters: str) -> PandasLikeSeries:
        return self.with_native(self.native.str.lstrip(characters))

    def strip_chars_end(self, characters: str) -> PandasLikeSeries:
        return self.with_native(self.native.str.rstrip(characters))

    def starts_with(self, prefix: PandasLikeSeries) -> PandasLikeSeries:
        _, prefix_native = align_and_extract_native(self.compliant, prefix)
        if not isinstance(prefix_native, str):
            msg = f"`.str.starts_with` only supports prefix values for {self.compliant._implementation} backend"
            raise TypeError(msg)
        return self._with_native_bool(self.native.str.startswith(prefix_native))

    def ends_with(self, suffix: PandasLikeSeries) -> PandasLikeSeries:
        _, suffix_native = align_and_extract_native(self.compliant, suffix)
        if not isinstance(suffix_native, str):
            msg = f"`.str.ends_with` only supports suffix values for {self.compliant._implementation} backend"
            raise TypeError(msg)
        return self._with_native_bool(self.native.str.endswith(suffix_native))

    def contains(self, pattern: PandasLikeSeries, *, literal: bool) -> PandasLikeSeries:
        _, pattern_native = align_and_extract_native(self.compliant, pattern)
        if not isinstance(pattern_native, str):
            msg = f"`.str.contains` only supports str pattern values for {self.compliant._implementation} backend"
            raise TypeError(msg)
        return self._with_native_bool(
            self.native.str.contains(pat=pattern_native, regex=not literal)
        )

    def slice(self, offset: int, length: int | None) -> PandasLikeSeries:
        stop = offset + length if length else None
        return self.with_native(self.native.str.slice(start=offset, stop=stop))

    def split(self, by: str) -> PandasLikeSeries:
        implementation = self.implementation
        if not implementation.is_cudf() and not is_dtype_pyarrow(self.native.dtype):
            msg = (
                "This operation requires a pyarrow-backed series. "
                "Please refer to https://narwhals-dev.github.io/narwhals/api-reference/narwhals/#narwhals.maybe_convert_dtypes "
                "and ensure you are using dtype_backend='pyarrow'. "
                "Additionally, make sure you have pandas version 1.5+ and pyarrow installed. "
            )
            raise TypeError(msg)
        return self.with_native(self.native.str.split(pat=by))

    def to_datetime(self, format: str | None) -> PandasLikeSeries:
        # If we know inputs are timezone-aware, we can pass `utc=True` for better performance.
        if format and any(x in format for x in ("%z", "Z")):
            return self.with_native(self._to_datetime(format, utc=True))
        result = self.with_native(self._to_datetime(format, utc=False))
        if (tz := getattr(result.dtype, "time_zone", None)) and tz != "UTC":
            return result.dt.convert_time_zone("UTC")
        return result

    def _to_datetime(self, format: str | None, *, utc: bool) -> Any:
        result = self.implementation.to_native_namespace().to_datetime(
            self.native, format=format, utc=utc
        )
        return (
            result.convert_dtypes(dtype_backend="pyarrow")
            if is_dtype_pyarrow(self.native.dtype)
            else result
        )

    def to_date(self, format: str | None) -> PandasLikeSeries:
        return self.to_datetime(format=format).dt.date()

    def to_time(self, format: str | None) -> PandasLikeSeries:
        time_dtype = self.version.dtypes.Time()
        return self.with_native(self._to_datetime(format, utc=False)).cast(time_dtype)

    def to_uppercase(self) -> PandasLikeSeries:
        return self.with_native(self.native.str.upper())

    def to_lowercase(self) -> PandasLikeSeries:
        return self.with_native(self.native.str.lower())

    def to_titlecase(self) -> PandasLikeSeries:
        return self.with_native(self.native.str.title())

    def zfill(self, width: int) -> PandasLikeSeries:
        return self.with_native(self.native.str.zfill(width))

    def pad_start(self, length: int, fill_char: str) -> PandasLikeSeries:
        return self.with_native(
            self.native.str.pad(width=length, fillchar=fill_char, side="left")
        )

    def pad_end(self, length: int, fill_char: str) -> PandasLikeSeries:
        return self.with_native(
            self.native.str.pad(width=length, fillchar=fill_char, side="right")
        )
