# ========= 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 asyncio
from collections.abc import Iterable
from dataclasses import dataclass
from enum import StrEnum

from mirage.core.timeutil import epoch_to_iso
from mirage.resource.base import BaseResource
from mirage.types import LINK_TARGET_KEY, FileStat, FileType, MountMode
from mirage.utils.path import glob_prefix_match, resolve_symlinks
from mirage.workspace.mount.mount import MountEntry
from mirage.workspace.mount.namespace.ram import RAMNamespaceStore
from mirage.workspace.mount.namespace.store import NamespaceStore, NodeFields
from mirage.workspace.mount.registry import MountRegistry


class NodeMetaKey(StrEnum):
    TARGET = "target"
    MTIME = "mtime"
    MODE = "mode"
    UID = "uid"
    GID = "gid"
    ATIME = "atime"
    OBSERVED_MTIME = "observed_mtime"


@dataclass(slots=True)
class NodeMeta:
    """Per-path namespace metadata.

    Two roles, distinguished by ``target``: a target-bearing entry is an
    authoritative symlink (the link exists only here); a target-less entry
    is a metadata overlay for a backend file whose backend has no native
    attribute slot. Attributes are stored, not enforced: mount mode does
    real access control.
    """

    target: str | None = None
    mtime: float | None = None
    mode: int | None = None
    uid: int | str | None = None
    gid: int | str | None = None
    atime: str | None = None
    # Time of the last content write observed through mirage. Unlike
    # `mtime` (an explicit touch, which wins over the backend), this is
    # a fallback that only surfaces when the backend reports no
    # modified time at all, so `find -mtime` works on mtime-less
    # backends for files written through mirage.
    observed_mtime: float | None = None

    def is_empty(self) -> bool:
        return all(getattr(self, key) is None for key in NodeMetaKey)

    def to_fields(self) -> NodeFields:
        return {
            str(key): value
            for key in NodeMetaKey if (value := getattr(self, key)) is not None
        }

    @classmethod
    def from_fields(cls, entry: NodeFields) -> "NodeMeta":
        target = entry.get(NodeMetaKey.TARGET)
        mtime = entry.get(NodeMetaKey.MTIME)
        mode = entry.get(NodeMetaKey.MODE)
        uid = entry.get(NodeMetaKey.UID)
        gid = entry.get(NodeMetaKey.GID)
        atime = entry.get(NodeMetaKey.ATIME)
        observed = entry.get(NodeMetaKey.OBSERVED_MTIME)
        return cls(
            target=target if isinstance(target, str) else None,
            mtime=float(mtime) if isinstance(mtime, (int, float)) else None,
            mode=mode if isinstance(mode, int) else None,
            uid=uid if isinstance(uid, (int, str)) else None,
            gid=gid if isinstance(gid, (int, str)) else None,
            atime=atime if isinstance(atime, str) else None,
            observed_mtime=(float(observed) if isinstance(
                observed, (int, float)) else None),
        )


def link_stat(name: str, meta: NodeMeta) -> FileStat:
    """Render a symlink node as a stat row.

    Size is the target string's byte length and mode is left unset so
    the formatter supplies 0777, which is what a real symlink inode
    reports (a link carries no permission bits of its own). Ownership
    and the stamps are not in that category: a link has a real uid/gid
    that ``chown -h`` writes and ``ls -l`` shows, and real times that
    ``touch -h`` and a guest's no-follow ``utime`` write, so all four
    ride through from the node. The
    target rides along in ``extra`` so every surface that has to name it
    (ls -l's ``name -> target``, file's "symbolic link to") reads one
    fact rather than querying the link table a second time.

    Args:
        name (str): name to report, usually the link's basename.
        meta (NodeMeta): the node entry; its target must not be None.
    """
    target = meta.target or ""
    return FileStat(
        name=name,
        size=len(target.encode("utf-8")),
        modified=epoch_to_iso(meta.mtime) if meta.mtime is not None else None,
        type=FileType.SYMLINK,
        uid=meta.uid,
        gid=meta.gid,
        atime=meta.atime,
        extra={LINK_TARGET_KEY: target},
    )


class Namespace:
    """Addressing authority: maps virtual paths to their mounts.

    Owns the mount registry and the per-path node-metadata table (symlinks
    plus the attribute overlay). Pure addressing: it resolves a virtual
    path to its mount and backend-relative path, following symlinks and
    crossing mounts. It holds no cache and performs no backend I/O. Op
    execution and caching live in the Dispatcher, which calls this layer
    to locate the mount.

    The in-memory table is the working copy (reads stay synchronous on
    the hot path); every mutation writes through to a ``NamespaceStore``
    (RAM by default, Redis for restart-durable namespaces), and the
    table hydrates from the store once on first use.

    Symlinks are stored verbatim as typed: the target string is kept exactly
    as the user wrote it (relative targets are resolved lazily against the
    link's own parent at resolution time), so ``readlink`` is GNU-faithful.
    """

    def __init__(self,
                 registry: MountRegistry,
                 store: NamespaceStore | None = None,
                 user: str | None = None) -> None:
        self._registry = registry
        self._store = store if store is not None else RAMNamespaceStore()
        self._nodes: dict[str, NodeMeta] = {}
        self._loaded = False
        self._load_lock = asyncio.Lock()
        self._claim = user
        self._user: str | None = None
        self._user_resolved = False

    @property
    def registry(self) -> MountRegistry:
        return self._registry

    @property
    def nodes(self) -> dict[str, NodeMeta]:
        return self._nodes

    @property
    def user(self) -> str | None:
        """The workspace user (whoami identity).

        Before store resolution the launch claim answers; after it, the
        resolved identity. None when no agent ever claimed the
        workspace, mirroring a uid with no passwd entry.
        """
        if self._user is not None:
            return self._user
        return self._claim

    async def _resolve_user(self) -> None:
        """Resolve the workspace user against the store, once.

        An explicit launch claim wins and writes through; without one the
        stored identity is adopted, so a runtime attaching to a shared
        store (e.g. Redis) inherits the workspace's whoami.
        """
        if self._user_resolved:
            return
        stored = await self._store.load_user()
        if self._claim is not None:
            self._user = self._claim
            if stored != self._claim:
                await self._store.set_user(self._claim)
        else:
            self._user = stored
        self._user_resolved = True

    async def ensure_loaded(self) -> None:
        """Hydrate the working copy from the store, once.

        A snapshot restore (``replace_nodes``) marks the table loaded, so
        snapshot state wins over whatever the store held before it.
        """
        if self._loaded and self._user_resolved:
            return
        async with self._load_lock:
            if not self._loaded:
                entries = await self._store.load()
                self._nodes = {
                    path: NodeMeta.from_fields(entry)
                    for path, entry in entries.items()
                }
                self._loaded = True
            await self._resolve_user()

    async def replace_nodes(self, entries: dict[str, NodeMeta]) -> None:
        self._nodes = dict(entries)
        self._loaded = True
        await self._resolve_user()
        await self._store.replace_all({
            path: meta.to_fields()
            for path, meta in entries.items()
        })

    async def close(self) -> None:
        await self._store.close()

    def symlink_targets(self) -> dict[str, str]:
        return {
            path: meta.target
            for path, meta in self._nodes.items() if meta.target is not None
        }

    def has_links(self) -> bool:
        return any(meta.target is not None for meta in self._nodes.values())

    def is_link(self, path: str) -> bool:
        meta = self._nodes.get(path)
        return meta is not None and meta.target is not None

    def readlink(self, path: str) -> str | None:
        meta = self._nodes.get(path)
        return meta.target if meta is not None else None

    async def symlink(self, link: str, target: str, mtime: float) -> None:
        meta = self._nodes.setdefault(link, NodeMeta())
        meta.target = target
        meta.mtime = mtime
        await self._store.set(link, meta.to_fields())

    def meta_for(self, path: str) -> NodeMeta | None:
        return self._nodes.get(path)

    async def set_attrs(
        self,
        path: str,
        *,
        mode: int | None = None,
        uid: int | str | None = None,
        gid: int | str | None = None,
        atime: str | None = None,
        mtime: float | None = None,
    ) -> None:
        """Write overlay attributes for a path (setattr fallback).

        Used for symlinks (which have no backend inode) and for backends
        without a native metadata slot. Only non-None fields are written.

        Args:
            path (str): absolute virtual path.
            mode (int | None): permission bits (e.g. 0o644).
            uid (int | str | None): owner id or name.
            gid (int | str | None): group id or name.
            atime (str | None): ISO access time.
            mtime (float | None): modification time (epoch seconds).
        """
        meta = self._nodes.setdefault(path, NodeMeta())
        if mode is not None:
            meta.mode = mode
        if uid is not None:
            meta.uid = uid
        if gid is not None:
            meta.gid = gid
        if atime is not None:
            meta.atime = atime
        if mtime is not None:
            meta.mtime = mtime
        await self._store.set(path, meta.to_fields())

    async def drop_attrs(self, path: str, fields: Iterable[str]) -> None:
        """Drop overlay fields that a backend has applied natively.

        A residual-free native setattr means the real inode now holds the
        requested value, so a stale overlay field would shadow it forever
        (chmod 000 then chmod 644 must not keep showing 000). Symlink
        entries keep their target.

        Args:
            path (str): absolute virtual path.
            fields (Iterable[str]): NodeMeta field names to drop.
        """
        meta = self._nodes.get(path)
        if meta is None:
            return
        for field in fields:
            if field != str(NodeMetaKey.TARGET):
                setattr(meta, field, None)
        if meta.is_empty():
            del self._nodes[path]
            await self._store.delete([path])
            return
        await self._store.set(path, meta.to_fields())

    async def drop_overlay(self, path: str) -> bool:
        """Drop a metadata overlay whose backend file is gone (orphan GC).

        When a backend reports a path deleted, its attribute overlay is
        orphaned and would otherwise linger forever. Symlink entries are
        authoritative (they have no backend file), so they are left intact;
        only a target-less overlay node is removed.

        Args:
            path (str): absolute virtual path reported gone by the backend.

        Returns:
            bool: True when an overlay node was dropped.
        """
        meta = self._nodes.get(path)
        if meta is None or meta.target is not None:
            return False
        del self._nodes[path]
        await self._store.delete([path])
        return True

    async def clear_times(self,
                          path: str,
                          observed: float | None = None) -> None:
        """Drop overlay times after a content write.

        write(2) refreshes mtime, so a stored overlay time would
        otherwise shadow the backend's fresh one forever. Permission and
        ownership survive writes; a symlink entry keeps its own times.
        Content writes (not removals) also record the observed write
        time, the fallback merge_overlay_stat surfaces when the backend
        reports no mtime of its own.

        Args:
            path (str): absolute virtual path that was written.
            observed (float | None): epoch seconds of a content write
                to record; None for removal ops.
        """
        meta = self._nodes.get(path)
        if meta is None and observed is not None:
            meta = self._nodes.setdefault(path, NodeMeta())
        if meta is None or meta.target is not None:
            return
        meta.mtime = None
        meta.atime = None
        # None (a removal) also drops any prior observed time, so a
        # deleted path's meta does not linger.
        meta.observed_mtime = observed
        if meta.is_empty():
            del self._nodes[path]
            await self._store.delete([path])
            return
        await self._store.set(path, meta.to_fields())

    async def unlink(self, path: str) -> bool:
        if path in self._nodes:
            del self._nodes[path]
            await self._store.delete([path])
            return True
        return False

    async def unlink_glob(self, pattern: str) -> int:
        """Drop node entries matching an unexpanded glob operand.

        ``rm`` receives the pattern verbatim (backend wrappers expand
        globs themselves), so the node table must match it here. Drops
        matched entries and everything under a matched directory.

        Args:
            pattern (str): absolute virtual glob pattern.

        Returns:
            int: number of entries dropped.
        """
        doomed = [
            path for path in self._nodes if glob_prefix_match(path, pattern)
        ]
        for path in doomed:
            del self._nodes[path]
        if doomed:
            await self._store.delete(doomed)
        return len(doomed)

    async def rename(self, src: str, dst: str) -> bool:
        meta = self._nodes.pop(src, None)
        if meta is None:
            return False
        self._nodes[dst] = meta
        await self._store.set(dst, meta.to_fields())
        await self._store.delete([src])
        return True

    def follow(self, path: str) -> str:
        """Return ``path`` with all symlink prefixes resolved.

        Identity when no link entries exist or nothing matches.

        Args:
            path (str): absolute virtual path.

        Raises:
            CycleError: when resolution exceeds the hop limit (ELOOP).
        """
        targets = self.symlink_targets()
        if not targets:
            return path
        return resolve_symlinks(path, targets)

    def link_stat_at(self, path: str) -> FileStat | None:
        """lstat a path: the link's own stat, or None when not a link.

        A symlink has no backend inode, so the node table is the only
        authority for it. Every no-follow stat surface (the ops facade,
        the dispatcher, FUSE) answers through here so they cannot
        disagree about what a link looks like.

        Args:
            path (str): absolute virtual path.
        """
        meta = self._nodes.get(path.rstrip("/") or path)
        if meta is None or meta.target is None:
            return None
        return link_stat(path.rstrip("/").rsplit("/", 1)[-1], meta)

    def link_stats_below(self, directory: str) -> list[tuple[str, FileStat]]:
        """Every link at any depth under a directory, with its own stat.

        A walker (``find``) needs the whole subtree, unlike a listing
        (``ls``) which wants one level.

        Args:
            directory (str): absolute virtual directory path.
        """
        base = directory.rstrip("/") + "/"
        return [(path, link_stat(path.rsplit("/", 1)[-1], meta))
                for path, meta in self._nodes.items()
                if meta.target is not None and path.startswith(base)]

    def link_stats_under(self, directory: str) -> list[FileStat]:
        """Stat rows for the links living directly under a directory.

        Links are namespace state and invisible to a backend readdir, so
        listing commands merge these rows into the backend's entries
        before sorting, the way the kernel returns them from one
        readdir.

        Args:
            directory (str): absolute virtual directory path.
        """
        return [
            link_stat(name, meta)
            for name, meta in self._links_under(directory)
        ]

    def link_names_under(self, directory: str) -> set[str]:
        """The names of the links living directly under a directory.

        What a readdir row's link mark needs, which is a name question
        rather than a stat one: the door already holds every entry's
        stat and has only to learn which of those names the node table
        owns.

        Resolves a link prefix first, because a listing does: a readdir
        of ``/data/alias`` is dispatched at ``/data/real`` and answers
        with that directory's entries, so the marks have to come from
        there too or every link inside an aliased directory reads as
        whatever its followed stat said. ``link_stats_under`` needs no
        such resolution: it is handed the path the router already
        rewrote.

        Args:
            directory (str): absolute virtual directory path.
        """
        return {name for name, _ in self._links_under(self.follow(directory))}

    def _links_under(self, directory: str) -> list[tuple[str, NodeMeta]]:
        """The links living directly under a directory, as (name, meta).

        Args:
            directory (str): absolute virtual directory path.
        """
        base = directory.rstrip("/") + "/"
        return [(path[len(base):], meta) for path, meta in self._nodes.items()
                if meta.target is not None and path.startswith(base)
                and "/" not in path[len(base):]]

    async def purge_under(self, directory: str) -> int:
        """Drop every node entry under a directory (``rm -r`` semantics).

        Args:
            directory (str): absolute virtual directory path being removed.

        Returns:
            int: number of entries dropped.
        """
        base = directory.rstrip("/") + "/"
        doomed = [path for path in self._nodes if path.startswith(base)]
        for path in doomed:
            del self._nodes[path]
        if doomed:
            await self._store.delete(doomed)
        return len(doomed)

    def resolve(self,
                path: str,
                *,
                follow: bool = True) -> tuple[BaseResource, str, MountMode]:
        """Map a virtual path to ``(resource, resource_path, mode)``.

        Args:
            path (str): virtual path to resolve.
            follow (bool): follow symlinks (the node table) before mapping
                the path to its mount.

        Raises:
            CycleError: when symlink resolution exceeds the hop limit (ELOOP).
        """
        if follow:
            path = self.follow(path)
        return self._registry.resolve(path)

    def mount_for(self, path: str) -> MountEntry:
        return self._registry.mount_for(path)

    def try_mount_for(self, path: str) -> MountEntry | None:
        return self._registry.try_mount_for(path)

    def is_mount_root(self, path: str) -> bool:
        return self._registry.is_mount_root(path)
