# ========= 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 typing import Any

from mirage.cache.file import io as cache_io
from mirage.cache.manager import CacheManager
from mirage.io import IOResult
from mirage.observe.record import OpRecord
from mirage.ops.config import NO_FOLLOW_OPS
from mirage.types import ConsistencyPolicy, FileStat, PathSpec
from mirage.utils.key_prefix import mount_key
from mirage.workspace.mount import MountEntry
from mirage.workspace.mount.namespace import Namespace
from mirage.workspace.mount.namespace.overlay import merge_overlay_stat
from mirage.workspace.reconcile import Reconciler
from mirage.workspace.session import assert_mount_allowed

_DISPATCH_READ_OPS = frozenset({"read", "read_bytes"})
_DISPATCH_WRITE_OPS = frozenset({
    "write", "write_bytes", "append", "unlink", "create", "truncate", "mkdir",
    "rmdir", "rename"
})


class Dispatcher:
    """Route a single VFS op to its mount and keep the file cache + index
    consistent.

    Owns the cache/IO coordination that used to live on Workspace: cache
    lookups for read-caching backends, post-write file-cache eviction,
    and parent index invalidation. Constructed with the namespace (for
    addressing), cache store, and consistency policy; holds no other
    workspace state. Drift checking stays on Workspace (it reads/writes
    snapshot-owned state), which guards its own dispatch wrapper before
    delegating here.
    """

    def __init__(self, namespace: Namespace, cache,
                 consistency: ConsistencyPolicy) -> None:
        self._namespace = namespace
        self._cache = cache
        self._reconciler = Reconciler(cache, namespace, consistency)

    @property
    def reconciler(self) -> Reconciler:
        return self._reconciler

    async def dispatch(self, op: str, path: PathSpec,
                       **kwargs: Any) -> tuple[Any, IOResult]:
        if op not in NO_FOLLOW_OPS:
            followed = self._namespace.follow(path.virtual)
            if followed != path.virtual:
                path = PathSpec.from_str_path(followed)
        mount = self._namespace.mount_for(path.virtual)
        assert_mount_allowed(mount.prefix)
        caches_reads = mount.resource.caches_reads

        if caches_reads and op in _DISPATCH_READ_OPS:
            cached = await self._cache.get(path.virtual)
            if cached is not None and await self._reconciler.may_serve_cached(
                    mount, path.virtual):
                return cached, IOResult(reads={path.virtual: cached})

        if op == "rename" and isinstance(kwargs.get("dst"), PathSpec):
            # Ops.rename addresses both endpoints against the source's
            # mount; mirror that here so the backend sees a
            # mount-relative destination.
            dst = kwargs["dst"]
            kwargs["dst"] = PathSpec(
                virtual=dst.virtual,
                directory=dst.virtual.rsplit("/", 1)[0] or "/",
                resource_path=mount_key(dst.virtual, mount.prefix.rstrip("/")),
            )
        try:
            result = await mount.execute_op(op, path.virtual, **kwargs)
        except FileNotFoundError:
            await self._reconciler.on_op_missing(op, path.virtual)
            raise
        if op == "stat" and isinstance(result, FileStat):
            result = merge_overlay_stat(self._namespace.meta_for(path.virtual),
                                        result)
        if op in _DISPATCH_WRITE_OPS:
            await self.invalidate_after_write(mount, path)
            if op == "rename" and isinstance(kwargs.get("dst"), PathSpec):
                await self.invalidate_after_write(mount, kwargs["dst"])
        return result, IOResult()

    async def stat(self, path: str) -> FileStat:
        scope = PathSpec(virtual=path,
                         directory=path,
                         resource_path="",
                         resolved=True)
        result, _ = await self.dispatch("stat", scope)
        return result

    async def readdir(self, path: str) -> list[str]:
        scope = PathSpec(virtual=path,
                         directory=path,
                         resource_path="",
                         resolved=False)
        raw, _ = await self.dispatch("readdir", scope)
        return raw

    async def apply_io(self,
                       io: IOResult,
                       records: list[OpRecord] | None = None) -> None:
        await cache_io.apply_io(self._cache,
                                io,
                                self.is_cacheable_path,
                                records=records)

    def is_cacheable_path(self, path: str) -> bool:
        try:
            mount = self._namespace.mount_for(path)
        except ValueError:
            return False
        return mount.resource.caches_reads

    async def invalidate_after_write_by_path(self, path: str) -> None:
        """Drop file-cache + stale parent index after a write to `path`.

        Single source of truth for post-write invalidation. Called from
        both `Workspace.dispatch()` and `Ops._call(write=True)` so a
        write through any code path sees the same invalidation rules:
        file cache is dropped only for read-caching mounts, and the
        parent directory index is dirtied for any mount that maintains
        an index. No-op for paths that resolve to no known mount.

        Args:
            path (str): absolute mount path that was written.
        """
        try:
            mount = self._namespace.mount_for(path)
        except ValueError:
            return
        spec = PathSpec.from_str_path(
            path, mount_key(path, mount.prefix.rstrip("/")))
        await self.invalidate_after_write(mount, spec)

    async def invalidate_after_write(self, mount: MountEntry,
                                     path: PathSpec) -> None:
        await self._namespace.clear_times(path.virtual)
        manager = mount.cache_manager
        if manager is None:
            manager = CacheManager(self._cache, mount.resource.index,
                                   mount.prefix, mount.resource.caches_reads)
        await manager.invalidate_after_write(path)
