"""
hegm._ast_scan -- AST-based detection of training loop patterns.

Scans Python source to find ``for X in range(...): ... for Y, Z in enumerate(...):``
patterns. Returns line numbers (or call-site info) so HeGM can selectively patch
only the epoch range() call, avoiding false positives.
"""

import ast
import os
from dataclasses import dataclass
from pathlib import Path
from typing import AbstractSet, List, Union

# Cache: filename -> tuple of TrainingLoopSite for training loop range() call sites
_scan_cache: dict = {}


@dataclass(frozen=True)
class TrainingLoopSite:
    """A detected training-loop range() call site."""

    lineno: int
    col_offset: int
    # Outer loop target name (e.g. "epoch")
    loop_var: str
    # Inner loop uses enumerate
    has_enumerate: bool


def _is_range_call(node: ast.AST) -> bool:
    """True if node is a call to range() with one argument."""
    if not isinstance(node, ast.Call):
        return False
    func = node.func
    if isinstance(func, ast.Name):
        return func.id == "range"
    if isinstance(func, ast.Attribute):
        return func.attr == "range"
    return False


def _is_enumerate_call(node: ast.AST) -> bool:
    """True if node is a call to enumerate()."""
    if not isinstance(node, ast.Call):
        return False
    func = node.func
    if isinstance(func, ast.Name):
        return func.id == "enumerate"
    if isinstance(func, ast.Attribute):
        return func.attr == "enumerate"
    return False


def _body_contains_enumerate_loop(nodes: list) -> bool:
    """Recursively check if any For loop uses enumerate() as its iter."""
    for node in ast.walk(ast.Module(body=nodes)):
        if isinstance(node, ast.For):
            if _is_enumerate_call(node.iter):
                return True
    return False


def _get_lineno(node: ast.AST) -> int:
    """Get line number, with fallback for older AST."""
    return getattr(node, "lineno", 0)


def _get_col_offset(node: ast.AST) -> int:
    """Get column offset, with fallback."""
    return getattr(node, "col_offset", 0)


def find_training_loop_sites(
    source_or_path: Union[str, bytes, os.PathLike],
) -> List[TrainingLoopSite]:
    """Parse Python source and find range() call sites that are training loops.

    A training loop is: ``for X in range(...):`` where the body contains
    ``for Y, Z in enumerate(...):``.

    Parameters
    ----------
    source_or_path : str | bytes | PathLike
        Python source code as string, or path to a .py file.

    Returns
    -------
    list of TrainingLoopSite
        Call sites (lineno, col_offset, loop_var) for each detected epoch range.

    Examples
    --------
    >>> sites = find_training_loop_sites('''
    ... for epoch in range(num_epochs):
    ...     for step, batch in enumerate(dataloader):
    ...         pass
    ... ''')
    >>> sites[0].lineno
    2
    >>> sites[0].loop_var
    'epoch'
    """
    if isinstance(source_or_path, (str, bytes)) and "\n" in str(source_or_path):
        source = source_or_path
        cache_key = None
    else:
        path = Path(source_or_path)
        if not path.exists():
            return []
        cache_key = str(path.resolve())
        if cache_key in _scan_cache:
            return list(_scan_cache[cache_key])
        try:
            source = path.read_text(encoding="utf-8", errors="replace")
        except (OSError, UnicodeDecodeError):
            return []

    try:
        tree = ast.parse(source)
    except SyntaxError:
        return []

    sites: List[TrainingLoopSite] = []

    for node in ast.walk(tree):
        if isinstance(node, ast.For):
            iter_node = node.iter
            if _is_range_call(iter_node) and len(getattr(iter_node, "args", [])) == 1:
                if _body_contains_enumerate_loop(node.body):
                    loop_var = (
                        node.target.id
                        if isinstance(node.target, ast.Name)
                        else getattr(node.target, "id", "?")
                    )
                    sites.append(
                        TrainingLoopSite(
                            lineno=_get_lineno(iter_node),
                            col_offset=_get_col_offset(iter_node),
                            loop_var=loop_var,
                            has_enumerate=True,
                        )
                    )

    if cache_key is not None and sites:
        _scan_cache[cache_key] = tuple(sites)

    return sites


def get_training_loop_linenos(
    source_or_path: Union[str, bytes, os.PathLike],
) -> AbstractSet[int]:
    """Return set of line numbers for training-loop range() calls.

    Convenience wrapper for use in the range() proxy.
    """
    sites = find_training_loop_sites(source_or_path)
    return frozenset(s.lineno for s in sites)


def is_training_loop_site(filename: str, lineno: int) -> bool:
    """Check if (filename, lineno) is a training-loop range() call site.

    Used at runtime: when range() is called, get caller's file and line,
    then call this to decide whether to patch.

    Returns False for non-file sources (e.g. <stdin>, <string>, exec).
    """
    if not filename or filename.startswith("<"):
        return False
    try:
        path = Path(filename).resolve()
    except (OSError, ValueError):
        return False
    if not path.exists():
        return False
    key = str(path)
    if key not in _scan_cache:
        find_training_loop_sites(path)
    cached = _scan_cache.get(key, ())
    return lineno in {s.lineno for s in cached}
