"""Operator dispatch: one selection mechanism for all op families.

A family registers an ordered list of ``ImplRecord`` rows (name, impl
object, capability ``Spec``, machine-level ``available``).  Resolution:
explicit/context selection (strict — raises when incapable) > ``ASTR_OPS``
env entry (soft — falls through) > first capable row > family fallback.
``Spec`` is a composable predicate over ``CallContext``; the rows are the
family's decision table, printable via ``explain``.

Records flagged ``faithful=False`` change numerics (e.g. fp8) and are only
reachable through an explicit selection, never the implicit chain.
"""

import contextvars
import logging
import os
import threading
from contextlib import contextmanager
from dataclasses import dataclass, field, replace
from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple

import torch

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class CallContext:
    """Axes a decision may depend on, snapshotted per call.

    ``raw``/``raw_kwargs`` keep the original call so unregistered handles
    can still be probed through their own capability method.
    """

    family: str
    fwd: Optional[str] = None
    dtype: Optional[torch.dtype] = None
    device_cuda: bool = False
    ndim: Optional[int] = None
    head_dim: Optional[int] = None
    has_cache: bool = False
    has_mask: bool = False
    grad_enabled: bool = True
    raw: Tuple[Any, ...] = ()
    raw_kwargs: Mapping[str, Any] = field(default_factory=dict)

    def evolve(self, **changes) -> "CallContext":
        return replace(self, **changes)


class Spec:
    """Composable, self-describing predicate over ``CallContext``."""

    __slots__ = ("_fn", "_desc")

    def __init__(self, fn: Callable[[CallContext], bool], desc: str):
        self._fn = fn
        self._desc = desc

    def matches(self, ctx: CallContext) -> bool:
        return bool(self._fn(ctx))

    @property
    def description(self) -> str:
        return self._desc

    def __and__(self, other: "Spec") -> "Spec":
        return Spec(
            lambda ctx: self._fn(ctx) and other._fn(ctx),
            f"({self._desc} and {other._desc})",
        )

    def __or__(self, other: "Spec") -> "Spec":
        return Spec(
            lambda ctx: self._fn(ctx) or other._fn(ctx),
            f"({self._desc} or {other._desc})",
        )

    def __invert__(self) -> "Spec":
        return Spec(lambda ctx: not self._fn(ctx), f"not({self._desc})")

    @classmethod
    def always(cls) -> "Spec":
        return cls(lambda ctx: True, "always")

    @classmethod
    def of(cls, fn: Callable[[CallContext], bool], desc: str) -> "Spec":
        return cls(fn, desc)

    @classmethod
    def dtype_in(cls, *dtypes: torch.dtype) -> "Spec":
        names = [str(getattr(d, "name", d)) for d in dtypes]
        return cls(
            lambda ctx: ctx.dtype in dtypes, f"dtype∈{{{','.join(names)}}}"
        )

    @classmethod
    def ndim(cls, n: int) -> "Spec":
        return cls(lambda ctx: ctx.ndim == n, f"ndim=={n}")

    @classmethod
    def head_dim_in(cls, *dims: int) -> "Spec":
        return cls(
            lambda ctx: ctx.head_dim in dims,
            f"head_dim∈{{{','.join(map(str, dims))}}}",
        )

    @classmethod
    def fwd_in(cls, *modes: Optional[str]) -> "Spec":
        names = [str(m) for m in modes]
        return cls(lambda ctx: ctx.fwd in modes, f"fwd∈{{{','.join(names)}}}")

    @classmethod
    def inference(cls) -> "Spec":
        return cls(lambda ctx: ctx.fwd is not None, "inference")

    @classmethod
    def training(cls) -> "Spec":
        return cls(lambda ctx: ctx.fwd is None, "training")

    @classmethod
    def no_grad(cls) -> "Spec":
        return cls(lambda ctx: not ctx.grad_enabled, "no_grad")

    @classmethod
    def has_cache(cls) -> "Spec":
        return cls(lambda ctx: ctx.has_cache, "has_cache")

    @classmethod
    def mask_free(cls) -> "Spec":
        return cls(lambda ctx: not ctx.has_mask, "mask_free")

    @classmethod
    def cuda_device(cls) -> "Spec":
        return cls(lambda ctx: ctx.device_cuda, "cuda_device")


@dataclass(frozen=True)
class ImplRecord:
    """One decision-table row: an implementation plus its capability."""

    family: str
    name: str
    obj: Any
    spec: Spec
    available: Callable[[], bool] = lambda: True
    priority: int = 100
    faithful: bool = True


@dataclass
class OpFamily:
    name: str
    provider: Callable[[], List[ImplRecord]]
    fallback: Callable[[], ImplRecord]


_FAMILIES: Dict[str, OpFamily] = {}
_ENV_ALIASES: Dict[str, str] = {}

_current_overrides: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
    "astrai_op_overrides", default={}
)

_env_lock = threading.Lock()
_env_cache: Dict[tuple, Optional[Dict[str, str]]] = {}
_warned: set = set()


def register_family(
    name: str,
    provider: Callable[[], List[ImplRecord]],
    fallback: Callable[[], ImplRecord],
) -> None:
    """Register (or replace) a family; ``provider`` is re-evaluated per
    resolution so availability changes (tests, late imports) are honored."""
    _FAMILIES[name] = OpFamily(name, provider, fallback)


def register_env_alias(family: str, varname: str) -> None:
    """Legacy single-value env var for a family (e.g. attention →
    ASTR_BACKEND); an ASTR_OPS entry wins when both are set."""
    _ENV_ALIASES[family] = varname


def _family(name: str) -> OpFamily:
    fam = _FAMILIES.get(name)
    if fam is None:
        raise KeyError(f"no operator family registered under {name!r}")
    return fam


def _warn_once(message: str) -> None:
    if message not in _warned:
        _warned.add(message)
        logger.warning(message)


def set_override(family: str, handle: Any) -> contextvars.Token:
    overrides = dict(_current_overrides.get())
    overrides[family] = handle
    return _current_overrides.set(overrides)


def reset_override(token: contextvars.Token) -> None:
    _current_overrides.reset(token)


def get_override(family: str) -> Optional[Any]:
    return _current_overrides.get().get(family)


@contextmanager
def op_backend(**handles: Any):
    """Select implementations per family for the enclosed scope::

        with op_backend(attention="torch_native", rotary="torch"):
            engine.generate(...)

    String handles are validated eagerly against the family's currently
    available implementations; object handles pass through unchecked.
    """
    for family, handle in handles.items():
        if isinstance(handle, str):
            fam = _FAMILIES.get(family)
            if fam is None:
                raise ValueError(f"unknown operator family {family!r}")
            if _record_for_handle(fam, handle) is None:
                raise ValueError(f"Unknown {family} implementation: {handle!r}")
    tokens = [set_override(f, h) for f, h in handles.items()]
    try:
        yield
    finally:
        for token in reversed(tokens):
            reset_override(token)


def env_overrides() -> Dict[str, str]:
    """Merged ASTR_OPS + legacy-alias selections (family or "profile").

    Cached per distinct env content; unknown families / malformed entries
    warn once and are dropped (soft override, never fatal).
    """
    with _env_lock:
        merged: Dict[str, str] = {}
        raw = os.environ.get("ASTR_OPS", "").strip()
        if raw:
            key = ("ASTR_OPS", raw)
            if key not in _env_cache:
                parsed: Dict[str, str] = {}
                for item in raw.split(","):
                    key_part, sep, value = item.strip().partition("=")
                    key_part, value = key_part.strip(), value.strip()
                    if not sep or not key_part or not value:
                        _warn_once(f"ASTR_OPS: ignoring malformed entry {item!r}")
                        continue
                    parsed[key_part] = value
                _env_cache[key] = parsed or None
            merged.update(_env_cache[key] or {})
        for fam, varname in _ENV_ALIASES.items():
            raw = os.environ.get(varname, "").strip()
            if raw:
                key = (varname, raw)
                if key not in _env_cache:
                    _env_cache[key] = {fam: raw.lower()}
                merged.setdefault(fam, _env_cache[key][fam])
        for fam in [f for f in merged if f not in _FAMILIES and f != "profile"]:
            _warn_once(f"ASTR_OPS: unknown operator family {fam!r}; dropping it")
            merged.pop(fam)
        return merged


def env_selection(family: str) -> Optional[str]:
    return env_overrides().get(family)


@dataclass(frozen=True)
class Resolution:
    record: ImplRecord
    origin: str


class ExplicitSelectionError(RuntimeError):
    """An explicitly selected implementation cannot handle the call."""


def _record_for_handle(fam: OpFamily, handle: Any) -> Optional[ImplRecord]:
    records = sorted(fam.provider(), key=lambda r: r.priority)
    if isinstance(handle, str):
        return next((r for r in records if r.name == handle), None)
    return next((r for r in records if r.obj is handle), None)


def _adhoc_record(family: str, handle: Any) -> ImplRecord:
    """Wrap an unregistered object; capability probes its own method."""
    supports = getattr(handle, "supports_call", None)
    if supports is not None:
        spec = Spec.of(
            lambda ctx: bool(supports(*ctx.raw, **ctx.raw_kwargs)),
            f"{type(handle).__name__}.supports_call",
        )
    else:
        spec = Spec.always()
    return ImplRecord(family, type(handle).__name__, handle, spec)


def resolve(
    family: str, ctx: CallContext, explicit: Optional[Any] = None
) -> Resolution:
    """Resolve one family for one call (explicit-strict / implicit-loose)."""
    fam = _family(family)

    handle: Optional[Any] = None
    origin = "chain"
    if explicit is not None:
        handle, origin = explicit, "explicit"
    elif get_override(family) is not None:
        handle, origin = get_override(family), "context"
    else:
        env_name = env_selection(family)
        if env_name is not None:
            handle, origin = env_name, "env"

    if handle is not None:
        record = _record_for_handle(fam, handle)
        if record is None and not isinstance(handle, str):
            record = _adhoc_record(family, handle)
        if record is None:
            if origin in ("explicit", "context"):
                raise ValueError(f"Unknown {family} implementation: {handle!r}")
            _warn_once(f"ASTR_OPS: {family}={handle!r} is not registered; ignoring")
        else:
            if record.available() and record.spec.matches(ctx):
                return Resolution(record, origin)
            if origin in ("explicit", "context"):
                raise ExplicitSelectionError(
                    f"Explicitly-set backend {type(record.obj).__name__} cannot "
                    f"handle this {family} call; required: {record.spec.description}"
                )

    if handle is None and env_overrides().get("profile") == "reference":
        return Resolution(fam.fallback(), "profile")

    for record in sorted(fam.provider(), key=lambda r: r.priority):
        if record.available() and record.faithful and record.spec.matches(ctx):
            return Resolution(record, "chain")
    return Resolution(fam.fallback(), "fallback")


def resolve_plan(ctxs: Mapping[str, CallContext]) -> Dict[str, Resolution]:
    """Resolve several families at once (one decision snapshot)."""
    return {family: resolve(family, ctx) for family, ctx in ctxs.items()}


def _describe_ctx(ctx: CallContext) -> str:
    parts = [f"fwd={ctx.fwd}", f"dtype={ctx.dtype}"]
    if ctx.ndim is not None:
        parts.append(f"ndim={ctx.ndim}")
    if ctx.head_dim is not None:
        parts.append(f"head_dim={ctx.head_dim}")
    parts += [f"cache={ctx.has_cache}", f"mask={ctx.has_mask}", f"grad={ctx.grad_enabled}"]
    return " ".join(parts)


def explain(family: str, ctx: CallContext, explicit: Optional[Any] = None) -> str:
    """Human-readable decision trace for one family."""
    fam = _family(family)
    records = sorted(fam.provider(), key=lambda r: r.priority)
    lines = [f"[{family}] {_describe_ctx(ctx)}"]
    for record in records:
        if not record.available():
            lines.append(f"  {record.name}: SKIP unavailable")
        elif not record.faithful:
            lines.append(f"  {record.name}: SKIP not faithful (explicit-only)")
        elif record.spec.matches(ctx):
            lines.append(f"  {record.name}: MATCH ({record.spec.description})")
        else:
            lines.append(f"  {record.name}: reject ({record.spec.description})")
    try:
        resolution = resolve(family, ctx, explicit)
        lines.append(f"  => {resolution.record.name} (origin={resolution.origin})")
    except (ExplicitSelectionError, ValueError) as exc:
        lines.append(f"  => ERROR: {exc}")
    return "\n".join(lines)


def explain_plan(ctxs: Mapping[str, CallContext]) -> str:
    return "\n".join(explain(family, ctx) for family, ctx in ctxs.items())