# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import inspect
from collections.abc import AsyncIterator, Callable
from functools import partial
from typing import Any

from mirage.accessor.base import Accessor
from mirage.cache.context import CacheInvalidator, active_cache_manager
from mirage.types import PathSpec, PolymorphicReadFn, ReadBytesFn, ReadStreamFn


async def _serve_stream(manager: CacheInvalidator | None,
                        start: Callable[[], AsyncIterator[bytes]],
                        path: PathSpec) -> AsyncIterator[bytes]:
    # `start` defers the backend call so a warm hit never opens a stream.
    if manager is not None:
        cached = await manager.cached_bytes(path)
        if cached is not None:
            yield cached
            return
    source = start()
    try:
        async for chunk in source:
            yield chunk
    finally:
        close = getattr(source, "aclose", None)
        if close is not None:
            await close()


def cache_aware_read_stream(raw: ReadStreamFn) -> ReadStreamFn:
    """Wrap a backend ``read_stream`` so warm reads serve cached bytes.

    The returned reader keeps the backend's ``(accessor, path, ...)``
    signature, so it is a drop-in for the raw op on ``CommandIO`` (the
    factory wraps it there once per backend). On a warm hit it yields
    the whole cached blob as one chunk; otherwise it streams from the
    backend. ``cached_bytes`` is a no-op (returns None) for local or
    non-caching mounts, so this is safe to apply uniformly.

    The wrapper is a ``def`` (not an ``async def``) that captures the
    active cache manager eagerly and returns the async generator: the
    manager must be read when the command calls the reader (inside the
    mount's cache-manager scope), not lazily when the stream drains, by
    which time that scope is gone.

    Args:
        raw (ReadStreamFn): the backend ``read_stream`` op.
    """

    def reader(accessor: Accessor | None, path: PathSpec, *args: Any,
               **kwargs: Any) -> AsyncIterator[bytes]:
        manager = active_cache_manager()
        return _serve_stream(manager,
                             partial(raw, accessor, path, *args, **kwargs),
                             path)

    return reader


def cache_aware_bound_stream(raw: ReadStreamFn) -> ReadStreamFn:
    """Path-first twin of :func:`cache_aware_read_stream`.

    For readers injected into the generics, which arrive with accessor
    and index already bound (``bound_op``) and are called as
    ``read(path)``.

    Args:
        raw (ReadStreamFn): a bound ``read_stream`` reader.
    """

    def reader(path: PathSpec, *args: Any,
               **kwargs: Any) -> AsyncIterator[bytes]:
        manager = active_cache_manager()
        return _serve_stream(manager, partial(raw, path, *args, **kwargs),
                             path)

    return reader


def cache_aware_read_bytes(raw: ReadBytesFn) -> ReadBytesFn:
    """Wrap a backend ``read_bytes`` so warm reads serve cached bytes.

    Drop-in for the raw ``(accessor, path, ...)`` op, same signature (the
    factory wraps ``CommandIO`` ops with it). Returns the cached bytes on
    a warm hit, else reads from the backend. No-op for local or
    non-caching mounts.

    Args:
        raw (ReadBytesFn): the backend ``read_bytes`` op.
    """

    async def reader(accessor: Accessor | None, path: PathSpec, *args: Any,
                     **kwargs: Any) -> bytes:
        manager = active_cache_manager()
        if manager is not None:
            cached = await manager.cached_bytes(path)
            if cached is not None:
                return cached
        return await raw(accessor, path, *args, **kwargs)

    return reader


def cache_aware_bound_bytes(raw: ReadBytesFn) -> ReadBytesFn:
    """Path-first twin of :func:`cache_aware_read_bytes`.

    For readers injected into the generics, which arrive with accessor
    and index already bound (``bound_op``) and are called as
    ``read(path)``.

    Args:
        raw (ReadBytesFn): a bound ``read_bytes`` reader.
    """

    async def reader(path: PathSpec, *args: Any, **kwargs: Any) -> bytes:
        manager = active_cache_manager()
        if manager is not None:
            cached = await manager.cached_bytes(path)
            if cached is not None:
                return cached
        return await raw(path, *args, **kwargs)

    return reader


def cache_aware_read(raw: PolymorphicReadFn) -> PolymorphicReadFn:
    """Wrap a polymorphic bound reader so warm reads serve cached bytes.

    For the ``read`` contract used by ``head_multi`` / ``tail_multi`` /
    wc ``format_multi``: the reader arrives with accessor and index
    already bound, is called as ``read(path)``, and may return bytes, an
    awaitable of bytes, or an async byte iterator. On a warm hit the
    wrapped reader returns the cached bytes; otherwise it calls the raw
    reader and returns whatever it produced unchanged, so the consumer's
    own ``isawaitable`` / ``ensure_stream`` normalization still applies.
    No-op for local or non-caching mounts.

    The active cache manager is captured **eagerly**, when this wrapper
    is applied, not when the wrapped reader is later called: a consumer
    that yields lazily (``head_multi``) is drained after the mount's
    cache-manager scope is gone, so reading the contextvar at drain time
    would always miss. Apply this wrapper inside the command's scope
    (which the consumers do) so the captured manager travels with the
    stream, mirroring :func:`cache_aware_read_stream`.

    Args:
        raw (PolymorphicReadFn): bound bytes / awaitable / stream reader.
    """
    manager = active_cache_manager()

    async def reader(path: PathSpec, *args: Any,
                     **kwargs: Any) -> bytes | AsyncIterator[bytes]:
        if manager is not None:
            cached = await manager.cached_bytes(path)
            if cached is not None:
                return cached
        result = raw(path, *args, **kwargs)
        if inspect.isawaitable(result):
            return await result
        return result

    return reader


async def cached_prefix_bytes(path: PathSpec, n: int | None) -> bytes | None:
    """Return the first ``n`` cached bytes of ``path`` when warm, else None.

    Lets a range-read fast path (e.g. ``head -c N``) serve from a fully
    cached file without a partial backend fetch. ``n=None`` returns the
    whole cached blob.

    Args:
        path (PathSpec): the path to look up.
        n (int | None): byte count, or None for the whole file.
    """
    manager = active_cache_manager()
    if manager is None:
        return None
    cached = await manager.cached_bytes(path)
    if cached is None:
        return None
    return cached if n is None else cached[:n]
