# ========= 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. =========

from collections.abc import Awaitable
from datetime import datetime, timezone
from typing import cast

try:
    from redis.asyncio import Redis
except ImportError as _err:
    raise ImportError("RedisIndexCacheStore requires the 'redis' extra. "
                      "Install with: pip install mirage-ai[redis]") from _err

from mirage.cache.index.config import (IndexEntry, ListResult, LookupResult,
                                       LookupStatus)
from mirage.cache.index.store import IndexCacheStore
from mirage.core.timeutil import to_iso_z
from mirage.utils.key_prefix import under_path

ENTRY_PREFIX = "mirage:idx:entry:"
CHILDREN_PREFIX = "mirage:idx:children:"


def _text(value: str | bytes) -> str:
    return value.decode() if isinstance(value, bytes) else value


def _glob_escape(value: str) -> str:
    """Escape redis MATCH metacharacters in a literal path.

    A path may legally contain ``*?[]``, and SCAN's pattern is a glob, so
    an unescaped path would match keys it does not name. The escaping is
    a narrowing optimization only; the caller still filters the results
    at a path boundary.

    Args:
        value (str): A literal path to embed in a MATCH pattern.
    """
    out: list[str] = []
    for char in value:
        if char in "*?[]\\":
            out.append("\\")
        out.append(char)
    return "".join(out)


class RedisIndexCacheStore(IndexCacheStore):
    """Redis-backed index cache for remote resource metadata.

    Stores IndexEntry objects as JSON strings and directory children as
    Redis lists. Directory TTL is managed via native Redis key expiration.
    All writes within set_dir are batched in a single pipeline for efficiency.

    Multiple stores can share one Redis server by using distinct key_prefix
    values (e.g. "gdrive:", "s3:"). The full key layout is::

        {key_prefix}mirage:idx:entry:{resource_path}     -> JSON string
        {key_prefix}mirage:idx:children:{resource_path}  -> Redis list

    Args:
        ttl (float): Default time-to-live in seconds for directory listings.
        url (str): Redis connection URL, used when *client* is not provided.
        client (Redis | None): Pre-existing async Redis client. When given,
            the store will not close it on ``close()``.
        key_prefix (str): Namespace prefix prepended to every Redis key,
            allowing multiple stores to coexist on the same server.
    """

    def __init__(
        self,
        ttl: float = 600,
        url: str = "redis://localhost:6379/0",
        client: Redis | None = None,
        key_prefix: str = "",
    ) -> None:
        super().__init__()
        self._ttl = ttl
        self._client = (client if client is not None else Redis.from_url(
            url, decode_responses=True))
        self._owns_client = client is None
        self._pending_seed: tuple[dict[str, IndexEntry], dict[str, list[str]],
                                  datetime] | None = None
        p = key_prefix or ""
        self._entry_prefix = f"{p}{ENTRY_PREFIX}"
        self._children_prefix = f"{p}{CHILDREN_PREFIX}"

    def _entry_key(self, resource_path: str) -> str:
        return f"{self._entry_prefix}{resource_path}"

    def _children_key(self, resource_path: str) -> str:
        return f"{self._children_prefix}{resource_path}"

    def seed(self, entries: dict[str, IndexEntry],
             children: dict[str, list[str]], expires_at: datetime) -> None:
        self._pending_seed = (dict(entries), {
            path: list(keys)
            for path, keys in children.items()
        }, expires_at)

    async def _flush_seed(self) -> None:
        pending = self._pending_seed
        if pending is None:
            return
        entries, children, expires_at = pending
        now = datetime.now(timezone.utc)
        now_iso = to_iso_z(now)
        ttl_seconds = max(1, int((expires_at - now).total_seconds()))
        pipe = self._client.pipeline()
        for resource_path, entry in entries.items():
            if not entry.index_time:
                entry = entry.model_copy(update={"index_time": now_iso})
            pipe.set(self._entry_key(resource_path), entry.model_dump_json())
        for resource_path, child_keys in children.items():
            key = self._children_key(resource_path)
            pipe.delete(key)
            if child_keys:
                pipe.rpush(key, *child_keys)
                pipe.expire(key, ttl_seconds)
        await pipe.execute()
        if self._pending_seed is pending:
            self._pending_seed = None

    async def get(self, resource_path: str) -> LookupResult:
        await self._flush_seed()
        raw = await self._client.get(self._entry_key(resource_path))
        if raw is None:
            return LookupResult(status=LookupStatus.NOT_FOUND)
        entry = IndexEntry.model_validate_json(raw)
        return LookupResult(entry=entry)

    async def put(self, resource_path: str, entry: IndexEntry) -> None:
        await self._flush_seed()
        if not entry.index_time:
            entry = entry.model_copy(
                update={"index_time": to_iso_z(datetime.now(timezone.utc))})
        await self._client.set(self._entry_key(resource_path),
                               entry.model_dump_json())

    async def list_dir(self, resource_path: str) -> ListResult:
        await self._flush_seed()
        key = self._children_key(resource_path)
        exists = await self._client.exists(key)
        if not exists:
            return ListResult(status=LookupStatus.NOT_FOUND)
        ttl_remaining = await self._client.ttl(key)
        if ttl_remaining == -2:
            return ListResult(status=LookupStatus.EXPIRED)
        raw = await cast("Awaitable[list[str | bytes]]",
                         self._client.lrange(key, 0, -1))
        return ListResult(entries=[_text(entry) for entry in raw])

    async def set_dir(
        self,
        resource_path: str,
        entries: list[tuple[str, IndexEntry]],
        expired_at: datetime | None = None,
    ) -> None:
        await self._flush_seed()
        now = datetime.now(timezone.utc)
        now_iso = to_iso_z(now)
        prefix = "/" if resource_path == "/" else resource_path + "/"

        pipe = self._client.pipeline()
        child_keys: list[str] = []
        for name, entry in entries:
            full_path = prefix + name
            if not entry.index_time:
                entry = entry.model_copy(update={"index_time": now_iso})
            pipe.set(self._entry_key(full_path), entry.model_dump_json())
            child_keys.append(full_path)

        children_key = self._children_key(resource_path)
        pipe.delete(children_key)
        if child_keys:
            pipe.rpush(children_key, *child_keys)

        if expired_at:
            ttl_seconds = max(1, int((expired_at - now).total_seconds()))
        else:
            ttl_seconds = max(1, int(self._ttl))
        pipe.expire(children_key, ttl_seconds)

        await pipe.execute()

    async def entries(self) -> dict[str, IndexEntry]:
        await self._flush_seed()
        entries: dict[str, IndexEntry] = {}
        cursor = 0
        while True:
            cursor, keys = await self._client.scan(
                cursor, match=f"{self._entry_prefix}*", count=500)
            for key in keys:
                key_text = _text(key)
                raw = await self._client.get(key)
                if raw is not None:
                    resource_path = key_text.removeprefix(self._entry_prefix)
                    entries[resource_path] = IndexEntry.model_validate_json(
                        raw)
            if cursor == 0:
                return entries

    async def invalidate_dir(self, resource_path: str) -> None:
        await self._flush_seed()
        children_key = f"{self._children_prefix}{resource_path}"
        child_paths = await cast("Awaitable[list[str | bytes]]",
                                 self._client.lrange(children_key, 0, -1))
        pipe = self._client.pipeline()
        for child in child_paths:
            pipe.delete(self._entry_key(_text(child)))
        pipe.delete(children_key)
        await pipe.execute()

    async def _scan_delete(self, prefix: str, resource_path: str) -> None:
        """Delete every key under ``prefix`` naming a path in the subtree.

        Args:
            prefix (str): Key namespace to scan (entries or children).
            resource_path (str): Mount-absolute root of the subtree.
        """
        pattern = f"{prefix}{_glob_escape(resource_path.rstrip('/'))}*"
        cursor = 0
        while True:
            cursor, keys = await self._client.scan(cursor,
                                                   match=pattern,
                                                   count=500)
            doomed = [
                key for key in keys
                if under_path(_text(key).removeprefix(prefix), resource_path)
            ]
            if doomed:
                await self._client.delete(*doomed)
            if cursor == 0:
                return

    async def invalidate_prefix(self, resource_path: str) -> None:
        await self._flush_seed()
        await self._scan_delete(self._entry_prefix, resource_path)
        await self._scan_delete(self._children_prefix, resource_path)

    async def invalidate(self) -> None:
        """Clear rather than expire, because redis cannot say "stale" here.

        The RAM store marks entries expired in place, so a later lookup
        answers EXPIRED and a backend whose index *is* its listing knows to
        refetch. A redis key carries a real TTL and an expired one is
        simply gone, so absent and stale read the same and this can only
        clear. The consequence, deliberately chosen: a github mount on a
        redis index answers ENOENT after a CLI write instead of refetching
        (``ls`` reports the mount root missing). That is a loud failure,
        not a wrong answer -- a no-op here would instead serve the
        pre-write tree as if it were current, and quietly wrong is the
        worse of the two. Closing this properly means an ``invalidated_at``
        marker key compared against each entry's ``index_time``, which
        needs no schema change and can ride the same round trip.
        """
        await self.clear()

    async def clear(self) -> None:
        self._pending_seed = None
        cursor = 0
        while True:
            cursor, keys = await self._client.scan(
                cursor, match=f"{self._entry_prefix}*", count=500)
            if keys:
                await self._client.delete(*keys)
            if cursor == 0:
                break
        cursor = 0
        while True:
            cursor, keys = await self._client.scan(
                cursor, match=f"{self._children_prefix}*", count=500)
            if keys:
                await self._client.delete(*keys)
            if cursor == 0:
                break

    async def close(self) -> None:
        if self._closed:
            return
        if self._owns_client:
            await self._client.aclose()
        self._pending_seed = None
        await super().close()
