"""
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.0"
__all__ = ["memo", "FileRef", "evict", "join", "status", "data_dir", "resolve",
           "MIN_OPS"]

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

# Minimum compute cost (in normalized ops) for a result to be shared on
# the iroh network. Functions below this threshold are stored locally only
# and a warning is emitted suggesting diskcache or joblib.Memory instead.
#
# Default: 600 billion ops -- roughly 15 seconds on a 10-year-old 8-core
# desktop (~40 GFLOPS). On a modern 64-core server this is ~1 second of
# wall-clock, but the *work* is identical. coren normalizes wall-clock time
# into machine-independent ops, so this threshold is universal.
#
# Override with IROHDS_MIN_OPS.
MIN_OPS = int(os.environ.get("IROHDS_MIN_OPS", 600_000_000_000))


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)


# ---------------------------------------------------------------------------
# coren: machine capability + compute normalization (required dependency)
# ---------------------------------------------------------------------------

from coren import FnCost, MachCap

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

_client: Optional[DaemonClient] = None
_client_lock = threading.Lock()
_machcap = None
_machcap_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


def _get_machcap() -> MachCap:
    """Lazy singleton MachCap. Benchmarks run once per process."""
    global _machcap
    if _machcap is not None:
        return _machcap
    with _machcap_lock:
        if _machcap is None:
            _machcap = MachCap.read(data_dir(), detect_net=True, run_bench=True)
        return _machcap


# ---------------------------------------------------------------------------
# Peak memory measurement (cross-platform)
# ---------------------------------------------------------------------------

def _measure_peak_mem(fn, args, kwargs):
    """Run fn, return (result, compute_ns, peak_mem_bytes)."""
    _resource = None
    try:
        import resource as _resource  # Unix only; ImportError on Windows
    except ImportError:
        pass

    rss_before = 0
    if _resource:
        try:
            rss_before = _resource.getrusage(_resource.RUSAGE_SELF).ru_maxrss * 1024
        except OSError:
            pass

    t0 = time.monotonic_ns()
    result = fn(*args, **kwargs)
    compute_ns = time.monotonic_ns() - t0

    peak_mem = 0
    if _resource:
        try:
            rss_after = _resource.getrusage(_resource.RUSAGE_SELF).ru_maxrss * 1024
            peak_mem = max(0, rss_after - rss_before)
        except OSError:
            pass

    return result, compute_ns, peak_mem


# ---------------------------------------------------------------------------
# Machine-independent normalization + verdict (all logic in coren)
# ---------------------------------------------------------------------------

def _normalize(compute_ns: int, peak_mem: int, result_bytes: int) -> dict:
    """Convert measurement to machine-independent FnCost via coren roofline."""
    cost = _get_machcap().normalize(compute_ns, peak_mem, result_bytes)
    return {"ops": cost.ops, "mem_bytes": cost.mem_bytes, "peak_mem": peak_mem}


def _should_fetch(probe: dict) -> bool:
    """Decide fetch vs compute using coren's verdict().

    Returns True (fetch) if ops is 0 -- meaning the producer predates
    the FnCost gossip extension, so we have no cost info to decide with.
    """
    ops = probe.get("ops", 0)
    if ops == 0:
        return True  # no cost metadata from producer, default to fetch

    cost = FnCost(ops, probe.get("mem_bytes", 0),
                  probe.get("peak_mem", 0), probe.get("result_bytes", 0))
    return _get_machcap().verdict(cost).should_fetch()


# ---------------------------------------------------------------------------
# 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. Decide: fetch from network or compute locally?
        # probe() asks the daemon if any peer has this + returns FnCost metadata.
        # _should_fetch() runs coren's verdict against this machine's capabilities.
        probe = client.probe(ns, key)
        if probe is not None and _should_fetch(probe):
            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

        # No peers have it, or coren says compute is faster: compute locally.
        result, compute_ns, peak_mem = _measure_peak_mem(fn, args, kwargs)

        # Handle FileRef
        if isinstance(result, FileRef):
            file_size = os.path.getsize(result.abspath) if os.path.exists(result.abspath) else 0
            norm = _normalize(compute_ns, peak_mem, file_size)
            resp = client.put_file(
                ns, key, result.abspath,
                compute_ns=compute_ns, tag=effective_tag, **norm)
            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, **norm)
        else:
            blob = pickle.dumps(result, protocol=5)
            norm = _normalize(compute_ns, peak_mem, len(blob))
            client.put(ns, key, blob,
                       compute_ns=compute_ns, tag=effective_tag, **norm)

        if norm.get("ops", 0) < MIN_OPS:
            import warnings
            warnings.warn(
                f"irohds: {fn.__qualname__} normalized to "
                f"{norm.get('ops', 0) / 1e9:.1f} G ops of work, below the "
                f"{MIN_OPS / 1e9:.0f} G ops network threshold. This result "
                f"is cached locally only and will NOT be shared on the iroh "
                f"network. For lightweight functions, consider diskcache or "
                f"joblib.Memory instead.",
                stacklevel=2,
            )

        _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)
