"""pytest plugin: make our yaml shapes win over subclasses that hardcode their
shape list inside an overridden `get_input_iter` (instead of `self.shapes`),
which `_shape_inject_plugin` can't reach (conv padding / pool / nll /
scaled_softmax families, ~17 ops).

At `pytest_collection_finish` (all Benchmark subclasses defined), redirect
`get_input_iter` ONLY on classes whose own implementation does not reference
`self.shapes` (the source-level guard that protects well-behaved BLAS/generic
classes). The redirect yields one input per yaml shape via the op's standard
`input_fn(shape, dtype, device)`; ops with no yaml entry or a non-standard
input_fn fall through unchanged. This also bounds compile cost vs the large
hardcoded lists.
"""

from __future__ import annotations

import inspect
import os
import sys

import yaml as _yaml


_YAML_CACHE: dict | None = None


def _load_yaml(path: str) -> dict:
    if not path or not os.path.isfile(path):
        return {}
    try:
        with open(path, "r") as f:
            return _yaml.safe_load(f) or {}
    except Exception as exc:
        from _term_style import tag
        print(f"{tag('[shape-iter-inject-plugin]')} failed to load {path}: {exc}",
              file=sys.stderr)
        return {}


def _reads_self_shapes(func) -> bool:
    """True if the method references self.shapes (already covered by
    _shape_inject_plugin, so don't redirect). Defaults True on introspection
    failure (safe: leave the method alone)."""
    try:
        return "self.shapes" in inspect.getsource(func)
    except (OSError, TypeError):
        return True


def _accepts_shape_dtype_device(fn) -> bool:
    """True if `fn` can be called as `fn(shape, dtype, device)` (the standard
    input_fn signature). Guards against blas-style `fn(b,m,n,k,dtype,device,...)`.
    Defaults False on introspection failure (safe: fall back to original)."""
    if not callable(fn):
        return False
    try:
        inspect.signature(fn).bind(None, None, None)
    except (TypeError, ValueError):
        return False
    return True


def _make_patched(original):
    def patched(self, dtype):
        entry = (_YAML_CACHE or {}).get(getattr(self, "op_name", None))
        shapes = entry.get("shapes") if isinstance(entry, dict) else None
        input_fn = getattr(self, "input_fn", None)
        if shapes and _accepts_shape_dtype_device(input_fn):
            for shape in shapes:
                yield from input_fn(tuple(shape), dtype, self.device)
        else:
            yield from original(self, dtype)
    return patched


def pytest_collection_finish(session):
    global _YAML_CACHE
    from benchmark import base as fg_base
    from benchmark.conftest import Config as fg_Config

    _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")

    def _subclasses(cls):
        for sub in cls.__subclasses__():
            yield sub
            yield from _subclasses(sub)

    patched = 0
    seen = set()
    for cls in _subclasses(fg_base.Benchmark):
        if cls in seen:
            continue
        seen.add(cls)
        own = cls.__dict__.get("get_input_iter")
        # Redirect only classes with their own get_input_iter that ignore self.shapes.
        if own is not None and not _reads_self_shapes(own):
            cls.get_input_iter = _make_patched(own)
            patched += 1

    from _term_style import tag
    print(f"{tag('[shape-iter-inject-plugin]')} get_input_iter redirected on {patched} "
          f"hardcoded-shape Benchmark classes; yaml shapes now win",
          file=sys.stderr, flush=True)