"""pytest plugin: make our shape yaml win over subclasses that override
`set_shapes()` with hardcoded shapes (~9% of upstream Benchmark subclasses).

Wraps `Benchmark.init_user_config` (runs right after `set_shapes()`): if the
op_name has a yaml entry, overwrite `self.shapes` with the yaml shapes. Ops with
no yaml entry keep their own shapes. An arity guard skips ops whose subclass
normalizes shape arity after reading yaml (e.g. BLAS (B,M,N,K)↔(M,N,K)).
"""

from __future__ import annotations

import os
import sys

import yaml as _yaml


_YAML_CACHE: dict | None = None  # filled lazily from Config.shape_file
_PATCHED = False


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:
        print(f"[shape-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
        return {}


def _normalize_shapes(raw):
    """Mirror base.py's `[tuple(shape) for shape in self.shapes]`, defensively
    handling scalar entries (`5` → `(5,)`)."""
    out = []
    for shape in raw:
        if isinstance(shape, (list, tuple)):
            out.append(tuple(shape))
        else:
            out.append((shape,))
    return out


def _patch_init_user_config():
    global _PATCHED
    if _PATCHED:
        return
    from benchmark import base as fg_base
    from benchmark.conftest import Config as fg_Config

    original = fg_base.Benchmark.init_user_config

    def _shape_arity(shape):
        if isinstance(shape, (list, tuple)):
            return len(shape)
        return 1

    def patched(self):
        original(self)
        global _YAML_CACHE
        if _YAML_CACHE is None:
            _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
        entry = _YAML_CACHE.get(getattr(self, "op_name", None))
        if not isinstance(entry, dict):
            return
        shapes = entry.get("shapes")
        if not shapes:
            return
        # Skip ops whose subclass normalizes shape arity after reading yaml (e.g.
        # BLAS (B,M,N,K)↔(M,N,K)): detected as both sides being homogeneous in
        # arity but differing. Heterogeneous subclass shapes are safe to override.
        if getattr(self, "shapes", None) and len(self.shapes) > 0:
            cur_arities = {_shape_arity(s) for s in self.shapes}
            yaml_arities = {_shape_arity(s) for s in shapes}
            cur_homogeneous = len(cur_arities) == 1
            yaml_homogeneous = len(yaml_arities) == 1
            if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities:
                print(
                    f"[shape-inject-plugin] skip override for op_name={self.op_name!r}: "
                    f"subclass produced homogeneous arity {next(iter(cur_arities))}, "
                    f"yaml has homogeneous arity {next(iter(yaml_arities))} "
                    f"— subclass normalization preserved",
                    file=sys.stderr,
                )
                return
        self.shapes = _normalize_shapes(shapes)
        if "shape_desc" in entry:
            self.shape_desc = entry["shape_desc"]

    fg_base.Benchmark.init_user_config = patched
    _PATCHED = True


def pytest_configure(config):
    _patch_init_user_config()
    print("[shape-inject-plugin] Benchmark.init_user_config patched: "
          "yaml shapes now win over hardcoded subclass set_shapes()",
          file=sys.stderr, flush=True)