#!/usr/bin/env python3
"""
Pre-commit hook to detect time.sleep() calls in test files missing @pytest.mark.slow.

Tests that call time.sleep() with a delay > 0.1 seconds should be marked with
@pytest.mark.slow so CI can skip them with ``-m "not slow"``.

Uses Python AST to find real ``time.sleep(>0.1)`` calls (ignores mocked/patched
ones) and checks whether the enclosing test function or class carries the
``@pytest.mark.slow`` decorator.

Opt-out: add ``# allow: unmarked-sleep`` on the time.sleep() line.
"""

import ast
import re
import sys

SUPPRESS_RE = re.compile(r"#\s*allow:\s*unmarked-sleep(?:\s|$)")

# Minimum sleep duration (seconds) to flag.  Tiny sleeps used for thread
# yielding (0.001-0.05 s) are not worth marking slow.
THRESHOLD = 0.1


class SleepMarkerChecker(ast.NodeVisitor):
    """AST visitor that flags ``time.sleep(>THRESHOLD)`` without ``@pytest.mark.slow``."""

    def __init__(self, filepath: str, lines: list[str]):
        self.filepath = filepath
        self.lines = lines
        self.issues: list[tuple[int, float, str]] = []
        self._class_stack: list[ast.ClassDef] = []
        self._func_stack: list[ast.FunctionDef | ast.AsyncFunctionDef] = []
        # Track import aliases for ``import time as X``
        self._time_aliases: set[str] = {"time"}

    # -- import tracking ---------------------------------------------------

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            if alias.name == "time":
                self._time_aliases.add(alias.asname or alias.name)
        self.generic_visit(node)

    # -- scope tracking ----------------------------------------------------

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self._class_stack.append(node)
        self.generic_visit(node)
        self._class_stack.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        self._func_stack.append(node)
        self.generic_visit(node)
        self._func_stack.pop()

    visit_AsyncFunctionDef = visit_FunctionDef

    # -- detection ---------------------------------------------------------

    def visit_Call(self, node: ast.Call) -> None:
        func = node.func
        # Match ``time.sleep()``, ``t.sleep()`` (any alias of ``import time``)
        if (
            isinstance(func, ast.Attribute)
            and func.attr == "sleep"
            and isinstance(func.value, ast.Name)
            and func.value.id in self._time_aliases
        ):
            self._check_sleep(node)
        self.generic_visit(node)

    def _check_sleep(self, node: ast.Call) -> None:
        if not node.args:
            return

        val = self._constant_value(node.args[0])
        if val is None or val <= THRESHOLD:
            return

        # Check for suppress comment on the call's lines (handles multi-line calls)
        start = node.lineno - 1
        end = getattr(node, "end_lineno", node.lineno) - 1
        for idx in range(start, min(end + 1, len(self.lines))):
            if SUPPRESS_RE.search(self.lines[idx]):
                return

        # Check if enclosing function or class has @pytest.mark.slow
        if self._func_stack and self._has_slow_marker(self._func_stack[-1]):
            return
        if self._class_stack and self._has_slow_marker(self._class_stack[-1]):
            return

        func_name = (
            self._func_stack[-1].name if self._func_stack else "<module>"
        )
        cls_name = self._class_stack[-1].name if self._class_stack else None
        qual = f"{cls_name}.{func_name}" if cls_name else func_name
        self.issues.append((node.lineno, val, qual))

    # -- helpers -----------------------------------------------------------

    @staticmethod
    def _constant_value(node: ast.expr) -> float | None:
        """Return numeric value of an AST node, or None if not a constant."""
        if isinstance(node, ast.Constant) and isinstance(
            node.value, (int, float)
        ):
            return float(node.value)
        # Handle negative numbers: ast.UnaryOp(op=USub, operand=Constant)
        if (
            isinstance(node, ast.UnaryOp)
            and isinstance(node.op, ast.USub)
            and isinstance(node.operand, ast.Constant)
            and isinstance(node.operand.value, (int, float))
        ):
            return -float(node.operand.value)
        return None

    @staticmethod
    def _has_slow_marker(
        node: ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef,
    ) -> bool:
        """True when node is decorated with ``@pytest.mark.slow``."""
        for dec in node.decorator_list:
            if _is_pytest_mark_slow(dec):
                return True
        return False


def _is_pytest_mark_slow(node: ast.expr) -> bool:
    """Check whether a decorator is exactly ``@pytest.mark.slow``."""
    # Unwrap call: @pytest.mark.slow(reason="...")
    if isinstance(node, ast.Call):
        return _is_pytest_mark_slow(node.func)
    # @pytest.mark.slow  →  Attribute(value=Attribute(value=Name(id='pytest'), attr='mark'), attr='slow')
    if (
        isinstance(node, ast.Attribute)
        and node.attr == "slow"
        and isinstance(node.value, ast.Attribute)
        and node.value.attr == "mark"
        and isinstance(node.value.value, ast.Name)
        and node.value.value.id == "pytest"
    ):
        return True
    return False


def check_file(filepath: str) -> list[tuple[int, float, str]]:
    """Return list of (lineno, sleep_value, qualified_name) for violations."""
    try:
        with open(filepath, encoding="utf-8") as f:
            source = f.read()
    except Exception as exc:
        print(f"WARNING: Cannot read {filepath}: {exc}", file=sys.stderr)
        return []

    try:
        tree = ast.parse(source, filename=filepath)
    except SyntaxError as exc:
        print(f"WARNING: Syntax error in {filepath}: {exc}", file=sys.stderr)
        return []

    lines = source.splitlines()
    checker = SleepMarkerChecker(filepath, lines)
    checker.visit(tree)
    return checker.issues


def main() -> int:
    exit_code = 0
    violations = 0

    for filepath in sys.argv[1:]:
        issues = check_file(filepath)
        for lineno, val, qual in issues:
            violations += 1
            print(
                f"{filepath}:{lineno}: time.sleep({val}) in {qual} "
                f"missing @pytest.mark.slow"
            )
            exit_code = 1

    if exit_code:
        print()
        print(
            "Hint: either add @pytest.mark.slow to the test function/class, "
            "or replace time.sleep() with freezegun time travel."
        )
        print(
            "  Suppress with `# allow: unmarked-sleep` if the sleep is intentional."
        )
        print(f"  {violations} violation(s) found.")

    return exit_code


if __name__ == "__main__":
    sys.exit(main())
