from __future__ import annotations

import tempfile
from pathlib import Path

from codeflash_python._model import FunctionParent, FunctionToOptimize
from codeflash_python.context.pipeline import get_code_optimization_context

project_root = Path(__file__).parent.parent.resolve()


class HelperClass:
    """Helper used by OptimizeMe."""

    def helper_method(self, a, b, c):
        """Return the sum of three values."""
        return a + b + c


def OptimizeMe(a, b, c):
    """Delegate to HelperClass.helper_method."""
    return HelperClass().helper_method(a, b, c)


def test_get_outside_method_helper() -> None:
    """Context extraction finds a helper class method used by a module-level function."""
    file_path = Path(__file__).resolve()

    function_to_optimize = FunctionToOptimize(
        function_name="OptimizeMe",
        file_path=file_path,
        parents=(),
        starting_line=None,
        ending_line=None,
    )
    code_context = get_code_optimization_context(
        function_to_optimize, file_path.parent
    )
    helper_fqns = {
        h.fully_qualified_name for h in code_context.helper_functions
    }
    assert "test_get_helper_code.HelperClass.helper_method" in helper_fqns


def test_flavio_typed_code_helper() -> None:
    """Context extraction resolves helpers through typed protocol classes."""
    code = '''

_P = ParamSpec("_P")
_KEY_T = TypeVar("_KEY_T")
_STORE_T = TypeVar("_STORE_T")
class AbstractCacheBackend(CacheBackend, Protocol[_KEY_T, _STORE_T]):
    """Interface for cache backends used by the persistent cache decorator."""

    def __init__(self) -> None: ...

    def hash_key(
        self,
        *,
        func: Callable[_P, Any],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> tuple[str, _KEY_T]: ...

    def encode(self, *, data: Any) -> _STORE_T:  # noqa: ANN401
        ...

    def decode(self, *, data: _STORE_T) -> Any:  # noqa: ANN401
        ...

    def get(self, *, key: tuple[str, _KEY_T]) -> tuple[datetime.datetime, _STORE_T] | None: ...

    def delete(self, *, key: tuple[str, _KEY_T]) -> None: ...

    def put(self, *, key: tuple[str, _KEY_T], data: _STORE_T) -> None: ...

    def get_cache_or_call(
        self,
        *,
        func: Callable[_P, Any],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
        lifespan: datetime.timedelta,
    ) -> Any:  # noqa: ANN401
        """
        Retrieve the cached results for a function call.

        Args:
        ----
            func (Callable[..., _R]): The function to retrieve cached results for.
            args (tuple[Any, ...]): The positional arguments passed to the function.
            kwargs (dict[str, Any]): The keyword arguments passed to the function.
            lifespan (datetime.timedelta): The maximum age of the cached results.

        Returns:
        -------
            _R: The cached results, if available.

        """
        if os.environ.get("NO_CACHE"):
            return func(*args, **kwargs)

        try:
            key = self.hash_key(func=func, args=args, kwargs=kwargs)
        except:  # noqa: E722
            # If we can't create a cache key, we should just call the function.
            logging.warning("Failed to hash cache key for function: %s", func)
            return func(*args, **kwargs)
        result_pair = self.get(key=key)

        if result_pair is not None:
            cached_time, result = result_pair
            if not os.environ.get("RE_CACHE") and (
                datetime.datetime.now() < (cached_time + lifespan)  # noqa: DTZ005
            ):
                try:
                    return self.decode(data=result)
                except CacheBackendDecodeError as e:
                    logging.warning("Failed to decode cache data: %s", e)
                    # If decoding fails we will treat this as a cache miss.
                    # This might happens if underlying class definition of the data changes.
            self.delete(key=key)
        result = func(*args, **kwargs)
        try:
            self.put(key=key, data=self.encode(data=result))
        except CacheBackendEncodeError as e:
            logging.warning("Failed to encode cache data: %s", e)
        # If encoding fails, we should still return the result.
        return result

_P = ParamSpec("_P")
_R = TypeVar("_R")
_CacheBackendT = TypeVar("_CacheBackendT", bound=CacheBackend)


class _PersistentCache(Generic[_P, _R, _CacheBackendT]):
    """
    A decorator class that provides persistent caching functionality for a function.

    Args:
    ----
        func (Callable[_P, _R]): The function to be decorated.
        duration (datetime.timedelta): The duration for which the cached results should be considered valid.
        backend (_backend): The backend storage for the cached results.

    Attributes:
    ----------
        __wrapped__ (Callable[_P, _R]): The wrapped function.
        __duration__ (datetime.timedelta): The duration for which the cached results should be considered valid.
        __backend__ (_backend): The backend storage for the cached results.

    """  # noqa: E501

    __wrapped__: Callable[_P, _R]
    __duration__: datetime.timedelta
    __backend__: _CacheBackendT

    def __init__(
        self,
        func: Callable[_P, _R],
        duration: datetime.timedelta,
    ) -> None:
        self.__wrapped__ = func
        self.__duration__ = duration
        self.__backend__ = AbstractCacheBackend()
        functools.update_wrapper(self, func)

    def cache_clear(self) -> None:
        """Clears the cache for the wrapped function."""
        self.__backend__.del_func_cache(func=self.__wrapped__)

    def no_cache_call(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
        """
        Calls the wrapped function without using the cache.

        Args:
        ----
            *args (_P.args): Positional arguments for the wrapped function.
            **kwargs (_P.kwargs): Keyword arguments for the wrapped function.

        Returns:
        -------
            _R: The result of the wrapped function.

        """
        return self.__wrapped__(*args, **kwargs)

    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
        """
        Calls the wrapped function, either using the cache or bypassing it based on environment variables.

        Args:
        ----
            *args (_P.args): Positional arguments for the wrapped function.
            **kwargs (_P.kwargs): Keyword arguments for the wrapped function.

        Returns:
        -------
            _R: The result of the wrapped function.

        """  # noqa: E501
        if "NO_CACHE" in os.environ:
            return self.__wrapped__(*args, **kwargs)
        os.makedirs(DEFAULT_CACHE_LOCATION, exist_ok=True)
        return self.__backend__.get_cache_or_call(
            func=self.__wrapped__,
            args=args,
            kwargs=kwargs,
            lifespan=self.__duration__,
        )
'''
    with tempfile.TemporaryDirectory() as tempdir:
        tempdir_path = Path(tempdir)
        file_path = (tempdir_path / "typed_code_helper.py").resolve()
        file_path.write_text(code, encoding="utf-8")
        project_root_path = tempdir_path.resolve()
        function_to_optimize = FunctionToOptimize(
            function_name="__call__",
            file_path=file_path,
            parents=(
                FunctionParent(name="_PersistentCache", type="ClassDef"),
            ),
            starting_line=None,
            ending_line=None,
        )
        code_context = get_code_optimization_context(
            function_to_optimize, project_root_path
        )
        helper_qns = {h.qualified_name for h in code_context.helper_functions}
        assert "AbstractCacheBackend.get_cache_or_call" in helper_qns
        # Verify the testgen context contains the key parts
        md = code_context.testgen_context.markdown
        assert "AbstractCacheBackend" in md
        assert "get_cache_or_call" in md
        assert "_PersistentCache" in md
        assert "__call__" in md
        # The context pipeline prunes cache_clear and no_cache_call since
        # they are not in the call chain of __call__
        assert "cache_clear" not in md
        assert "no_cache_call" not in md


def test_bubble_sort_deps() -> None:
    """Context extraction follows cross-file dependencies for bubble_sort_deps."""
    file_path = (
        Path(__file__) / ".." / "code_to_optimize" / "bubble_sort_deps.py"
    ).resolve()

    function_to_optimize = FunctionToOptimize(
        function_name="sorter_deps",
        file_path=file_path,
        parents=(),
        starting_line=None,
        ending_line=None,
    )
    proj_root = file_path.parent.parent.resolve()
    code_context = get_code_optimization_context(
        function_to_optimize, proj_root
    )
    dep1_path = Path("code_to_optimize/bubble_sort_dep1_helper.py").as_posix()
    dep2_path = Path("code_to_optimize/bubble_sort_dep2_swap.py").as_posix()
    deps_path = Path("code_to_optimize/bubble_sort_deps.py").as_posix()
    expected = (
        f"```python:{dep1_path}\n"
        "def dep1_comparer(arr, j: int) -> bool:\n"
        "    return arr[j] > arr[j + 1]\n"
        "```\n"
        f"```python:{dep2_path}\n"
        "def dep2_swap(arr, j):\n"
        "    temp = arr[j]\n"
        "    arr[j] = arr[j + 1]\n"
        "    arr[j + 1] = temp\n"
        "```\n"
        f"```python:{deps_path}\n"
        "from code_to_optimize.bubble_sort_dep1_helper import dep1_comparer\n"
        "from code_to_optimize.bubble_sort_dep2_swap import dep2_swap\n"
        "\n"
        "\n"
        "def sorter_deps(arr):\n"
        "    for i in range(len(arr)):\n"
        "        for j in range(len(arr) - 1):\n"
        "            if dep1_comparer(arr, j):\n"
        "                dep2_swap(arr, j)\n"
        "    return arr\n"
        "```"
    )
    assert code_context.testgen_context.markdown == expected
    helper_fqns = {
        h.fully_qualified_name for h in code_context.helper_functions
    }
    assert (
        "code_to_optimize.bubble_sort_dep1_helper.dep1_comparer" in helper_fqns
    )
    assert "code_to_optimize.bubble_sort_dep2_swap.dep2_swap" in helper_fqns
