"""
irohds: Decentralized function memoization over iroh P2P.

    import irohds

    @irohds.memo
    def f(a, b):
        time.sleep(30)
        return a + b

    f(10, 20)  # 30s first time, instant on any peer after
"""

import functools
import os
import pickle
import random
import time
import threading
from typing import Optional

from irohds._client import DaemonClient
from irohds._daemon import ensure_daemon
from irohds._hasher import make_cache_key

__version__ = "0.3.12"
__all__ = ["memo", "FileRef", "evict", "join", "status", "data_dir", "resolve"]

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


def data_dir() -> str:
    """Root directory for FileRef resolution. Override with IROHDS_DATA_DIR."""
    return os.environ.get("IROHDS_DATA_DIR", os.getcwd())


def resolve(rel_path: str) -> str:
    """Resolve a relative path against data_dir."""
    if os.path.isabs(rel_path):
        return rel_path
    return os.path.join(data_dir(), rel_path)


# ---------------------------------------------------------------------------
# Lazy singletons (thread-safe)
# ---------------------------------------------------------------------------

_client: Optional[DaemonClient] = None
_client_lock = threading.Lock()


def _get_client() -> DaemonClient:
    global _client
    if _client is not None:
        return _client
    with _client_lock:
        if _client is None:
            addr = ensure_daemon()
            _client = DaemonClient(addr)
        return _client


# ---------------------------------------------------------------------------
# FileRef: content-addressed handle for large files
# ---------------------------------------------------------------------------

def _fileref_dir(cache_key: str) -> str:
    """Unique directory for materialized FileRef files, keyed by cache entry."""
    return os.path.join(data_dir(), ".irohds", "files", cache_key[:16])


class FileRef:
    """A content-addressed reference to a file.

    Paths are stored relative to ``irohds.data_dir()`` so the same file
    produces the same cache key regardless of absolute mount point.

    Usage::

        @irohds.memo
        def train_model(dataset, epochs=10):
            ...
            torch.save(model, irohds.resolve("model.pth"))
            return irohds.FileRef("model.pth")

        ref = train_model("cifar10", epochs=50)
        model = torch.load(ref.path)
    """

    __slots__ = ("rel_path", "blob_hash", "size")

    def __init__(self, path: str, blob_hash: str = "", size: int = 0):
        self.rel_path = self._to_relative(path)
        self.blob_hash = blob_hash
        self.size = size

    @staticmethod
    def _to_relative(path: str) -> str:
        if not os.path.isabs(path):
            return path
        try:
            return os.path.relpath(path, data_dir())
        except ValueError:
            return path

    @property
    def abspath(self) -> str:
        return os.path.join(data_dir(), self.rel_path)

    @property
    def path(self) -> str:
        return self.abspath

    def _materialized_path(self, cache_key: str) -> str:
        return os.path.join(_fileref_dir(cache_key), os.path.basename(self.rel_path))

    def __repr__(self):
        if self.blob_hash:
            return f"FileRef({self.rel_path!r}, hash={self.blob_hash[:16]}..., size={self.size})"
        return f"FileRef({self.rel_path!r})"

    def __eq__(self, other):
        if not isinstance(other, FileRef):
            return NotImplemented
        return self.rel_path == other.rel_path and self.blob_hash == other.blob_hash

    def __getstate__(self):
        return (self.rel_path, self.blob_hash, self.size)

    def __setstate__(self, state):
        self.rel_path, self.blob_hash, self.size = state


def _materialize_fileref(result, key, client, ns):
    """Download and localize a FileRef if needed. Mutates result in place."""
    if not (isinstance(result, FileRef) and result.blob_hash):
        return
    target = result._materialized_path(key)
    if not os.path.exists(target):
        os.makedirs(os.path.dirname(target), exist_ok=True)
        size = client.get_file(ns, result.blob_hash, target)
        if size is not None:
            result.size = size
    result.rel_path = os.path.relpath(target, data_dir())


# ---------------------------------------------------------------------------
# @irohds.memo decorator
# ---------------------------------------------------------------------------

def _default_tag(fn) -> str:
    """module.qualname for human-readable eviction."""
    mod = getattr(fn, "__module__", "") or ""
    qual = getattr(fn, "__qualname__", "") or getattr(fn, "__name__", "")
    return f"{mod}.{qual}" if mod else qual


def memo(fn=None, *, ns="global", ignore=None, tag="", verify=0.0, max_depth=5):
    """Decentralized memoization decorator.

    Args:
        ns:        Namespace / gossip topic. "global" shares with everyone.
        ignore:    Parameter names to exclude from cache key.
        tag:       Tag prefix for cache keys. Defaults to module.qualname.
        verify:    Probability (0-1) of re-computing on hit to verify.
        max_depth: Transitive call graph depth for AST hashing.
    """
    if fn is None:
        return lambda f: memo(f, ns=ns, ignore=ignore, tag=tag,
                              verify=verify, max_depth=max_depth)

    effective_tag = tag or _default_tag(fn)

    # In-process result cache: {key: result}
    # Eliminates IPC entirely for repeated calls within the same process.
    # Results are immutable (it's a memoizer), so this is always correct.
    _cache = {}

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        key = make_cache_key(fn, args, kwargs,
                             ignore=ignore, tag=effective_tag,
                             max_depth=max_depth)

        # Fast path: in-process cache (no IPC, ~0.1us)
        if key in _cache:
            return _cache[key]

        client = _get_client()

        # Single IPC round-trip: check local cache in daemon
        data = client.check(ns, key)
        if data is not None:
            result = pickle.loads(data)

            # Materialize FileRef if needed
            _materialize_fileref(result, key, client, ns)

            if verify > 0 and random.random() < verify:
                actual = fn(*args, **kwargs)
                if pickle.dumps(result, protocol=5) != pickle.dumps(actual, protocol=5):
                    import warnings
                    warnings.warn(
                        f"irohds: verification mismatch for {fn.__qualname__}"
                        f" with key {key[:16]}...",
                        RuntimeWarning,
                    )

            _cache[key] = result
            return result

        # Local miss. Try fetching from network.
        data = client.get(ns, key)
        if data is not None:
            result = pickle.loads(data)

            _materialize_fileref(result, key, client, ns)

            _cache[key] = result
            return result

        # Network miss: compute locally.
        t0 = time.monotonic_ns()
        result = fn(*args, **kwargs)
        compute_ns = time.monotonic_ns() - t0

        # Handle FileRef
        if isinstance(result, FileRef):
            resp = client.put_file(
                ns, key, result.abspath,
                compute_ns=compute_ns, tag=effective_tag)
            result.blob_hash = resp.get("hash", "")
            result.size = resp.get("size", 0)
            blob = pickle.dumps(result, protocol=5)
            client.put(ns, key, blob,
                       compute_ns=compute_ns, tag=effective_tag)
        else:
            blob = pickle.dumps(result, protocol=5)
            client.put(ns, key, blob,
                       compute_ns=compute_ns, tag=effective_tag)

        _cache[key] = result
        return result

    wrapper.cache_key = lambda *a, **kw: make_cache_key(
        fn, a, kw, ignore=ignore, tag=effective_tag, max_depth=max_depth,
    )
    wrapper.uncached = fn
    wrapper.cache_clear = _cache.clear
    return wrapper


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def evict(tag: str, ns: str = "global"):
    """Evict all entries matching a tag prefix."""
    return _get_client().evict(ns, tag)


def join(ns: str = "global"):
    """Pre-join a namespace and start peer discovery."""
    _get_client().join(ns)


def status(ns: str = "global") -> dict:
    """Get daemon status for a namespace."""
    return _get_client().status(ns)
