"""Test code validation — split, validate individually, reassemble."""

from __future__ import annotations

import ast
import logging
import re

log = logging.getLogger(__name__)


class CodeValidationError(Exception):
    """
    Raised when generated test code cannot be validated.
    """

    def __init__(
        self,
        message: str,
        *,
        initial_code: str | None = None,
        fixed_code: str | None = None,
        final_code: str | None = None,
        lines_removed: int = 0,
        validation_error: str | None = None,
    ) -> None:
        """
        Store validation context alongside the error message.
        """
        super().__init__(message)
        self.initial_code = initial_code
        self.fixed_code = fixed_code
        self.final_code = final_code
        self.lines_removed = lines_removed
        self.validation_error = validation_error

    def to_debug_dict(self) -> dict[str, str | int | None]:
        """
        Convert to a dictionary for API responses.
        """
        return {
            "stage": "code_validation",
            "initial_code": self.initial_code,
            "fixed_code": self.fixed_code,
            "final_code": self.final_code,
            "lines_removed": self.lines_removed,
            "validation_error": self.validation_error,
        }


def _is_test_function_name(name: str) -> bool:
    """
    Check if *name* follows the test function naming convention.
    """
    return name.startswith("test_")


def _is_test_class_name(name: str) -> bool:
    """
    Check if *name* follows the test class naming convention.
    """
    return name.startswith("Test") or name.endswith(("Test", "TestCase"))


def _inherits_from_testcase(node: ast.ClassDef) -> bool:
    """
    Check if an AST class inherits from TestCase.
    """
    return any(
        (isinstance(base, ast.Name) and base.id == "TestCase")
        or (isinstance(base, ast.Attribute) and base.attr == "TestCase")
        for base in node.bases
    )


def _is_test_function(node: ast.stmt) -> bool:
    """
    Check if *node* is a test function.
    """
    if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
        return _is_test_function_name(node.name)
    return False


def _is_test_class(node: ast.stmt) -> bool:
    """
    Check if *node* is a test class.
    """
    if isinstance(node, ast.ClassDef):
        return _is_test_class_name(node.name) or _inherits_from_testcase(node)
    return False


def has_test_functions(tree: ast.Module) -> bool:
    """
    Check if AST contains pytest or unittest test functions.
    """

    class _Visitor(ast.NodeVisitor):
        """AST visitor that raises StopIteration on test functions."""

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            """Check sync function definitions."""
            if _is_test_function_name(node.name):
                raise StopIteration
            self.generic_visit(node)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            """Check async function definitions."""
            if _is_test_function_name(node.name):
                raise StopIteration
            self.generic_visit(node)

    try:
        _Visitor().visit(tree)
    except StopIteration:
        return True
    return False


_TEST_FUNCTION_PATTERN = re.compile(
    r"^[ \t]*((?:async\s+)?def\s+test_\w+\s*\([^)]*\)"
    r"\s*(?:->\s*[^:]+)?:)",
    re.MULTILINE,
)

_TEST_CLASS_PATTERN = re.compile(
    r"^(class\s+(?:Test\w*)\s*(?:\([^)]*\))?\s*:"
    r"|class\s+\w+\s*\(\s*(?:\w+\.)?TestCase\s*\)\s*:)",
    re.MULTILINE,
)


def _get_indent_level(line: str) -> int:
    """
    Get the number of leading whitespace characters.
    """
    return len(line) - len(line.lstrip())


def _find_decorator_start_line(lines: list[str], def_line: int) -> int:
    """
    Find the first decorator line preceding a definition.
    """
    if def_line <= 0:
        return def_line

    start_line = def_line
    i = def_line - 1
    max_look_back = 20

    while i >= 0 and (def_line - i) <= max_look_back:
        line = lines[i].strip()
        if line.startswith("@"):
            start_line = i
            i -= 1
        elif not line or line.startswith("#"):
            i -= 1
        else:
            break

    return start_line


def _split_code_with_ast(
    code: str,
    python_version: tuple[int, int],
) -> tuple[str, list[str]] | None:
    """
    Split code into preamble + test functions using AST.
    """
    lines = code.splitlines(keepends=True)

    try:
        tree = ast.parse(code, feature_version=python_version)
    except SyntaxError:
        return None

    preamble_lines: list[str] = []
    test_functions: list[str] = []
    last_preamble_end = 0

    for node in tree.body:
        if _is_test_function(node) or _is_test_class(node):
            if (
                isinstance(
                    node,
                    ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef,
                )
                and node.decorator_list
            ):
                start_line = node.decorator_list[0].lineno - 1
            else:
                start_line = node.lineno - 1
            end_line = node.end_lineno if node.end_lineno else start_line + 1

            if start_line > last_preamble_end:
                gap_lines = lines[last_preamble_end:start_line]
                if any(
                    line.strip() and not line.strip().startswith("#")
                    for line in gap_lines
                ):
                    preamble_lines.extend(gap_lines)  # pragma: no cover

            test_code = "".join(lines[start_line:end_line])
            test_functions.append(test_code)
            last_preamble_end = end_line
        else:
            if (
                isinstance(
                    node,
                    ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef,
                )
                and node.decorator_list
            ):
                start_line = node.decorator_list[0].lineno - 1
            else:
                start_line = node.lineno - 1
            end_line = node.end_lineno if node.end_lineno else start_line + 1

            if start_line > last_preamble_end:
                preamble_lines.extend(lines[last_preamble_end:start_line])

            preamble_lines.extend(lines[start_line:end_line])
            last_preamble_end = end_line

    preamble = "".join(preamble_lines).rstrip() + "\n"
    return preamble, test_functions


def _split_code_with_regex(
    code: str,
) -> tuple[str, list[str]]:
    """
    Split code using regex when AST parsing fails.
    """
    lines = code.splitlines(keepends=True)

    class_matches = list(_TEST_CLASS_PATTERN.finditer(code))

    top_level_test_funcs: list[re.Match[str]] = []
    for m in _TEST_FUNCTION_PATTERN.finditer(code):
        line_start = code.rfind("\n", 0, m.start()) + 1
        line_content = code[line_start : m.end()]
        if _get_indent_level(line_content) == 0:
            top_level_test_funcs.append(m)

    if not class_matches and not top_level_test_funcs:
        return code, []

    test_items: list[tuple[int, int, str]] = []

    for m in class_matches:
        line_start = code.rfind("\n", 0, m.start()) + 1
        def_line = code[:line_start].count("\n")
        start_line = _find_decorator_start_line(lines, def_line)
        test_items.append((m.start(), start_line, "class"))

    for m in top_level_test_funcs:
        line_start = code.rfind("\n", 0, m.start()) + 1
        def_line = code[:line_start].count("\n")
        start_line = _find_decorator_start_line(lines, def_line)
        test_items.append((m.start(), start_line, "func"))

    test_items.sort(key=lambda x: x[0])

    test_boundaries: list[tuple[int, int]] = []
    for i, (_start_pos, start_line, _item_type) in enumerate(test_items):
        end_line = (
            test_items[i + 1][1] if i + 1 < len(test_items) else len(lines)
        )
        test_boundaries.append((start_line, end_line))

    first_test_line = test_boundaries[0][0] if test_boundaries else len(lines)
    preamble_lines = lines[:first_test_line]
    preamble = "".join(preamble_lines).rstrip() + "\n"

    test_functions: list[str] = []
    for start_line, end_line in test_boundaries:
        test_code = "".join(lines[start_line:end_line]).rstrip()
        test_lines = test_code.splitlines(keepends=True)
        while test_lines and not test_lines[-1].strip():  # pragma: no cover
            test_lines.pop()
        if test_lines:
            test_functions.append("".join(test_lines).rstrip())

    return preamble, test_functions


def _split_code_into_parts(
    code: str,
    python_version: tuple[int, int],
) -> tuple[str, list[str]]:
    """
    Split code into preamble and individual test functions.
    """
    result = _split_code_with_ast(code, python_version)
    if result is not None:
        return result

    log.debug("AST parsing failed, using regex-based splitting")
    return _split_code_with_regex(code)


def _repair_preamble(
    preamble: str,
    python_version: tuple[int, int],
) -> str:
    """
    Attempt to fix a preamble with syntax errors.
    """
    lines = preamble.splitlines(keepends=True)
    max_attempts = min(len(lines), 15)

    for _ in range(max_attempts):
        current = "".join(lines).strip()
        if not current:
            return ""
        try:
            ast.parse(current, feature_version=python_version)
            return current
        except SyntaxError as e:
            if e.lineno is None:
                return ""
            error_idx = e.lineno - 1
            if 0 <= error_idx < len(lines):
                lines.pop(error_idx)
            else:
                return ""

    result = "".join(lines).strip()
    try:
        ast.parse(result, feature_version=python_version)
        return result
    except SyntaxError:
        return ""


def _validate_tests_individually(
    code: str,
    python_version: tuple[int, int],
) -> tuple[str, int]:
    """
    Validate each test function individually.
    """
    preamble, test_functions = _split_code_into_parts(code, python_version)

    if not test_functions:
        log.debug("No test functions found during splitting")
        return code, 0

    valid_tests: list[str] = []
    removed_count = 0
    preamble_stripped = preamble.strip()

    if preamble_stripped:
        try:
            ast.parse(preamble_stripped, feature_version=python_version)
        except SyntaxError:
            log.warning("Preamble has syntax error, attempting repair")
            preamble_stripped = _repair_preamble(
                preamble_stripped, python_version
            )

    for test_code in test_functions:
        combined = (
            preamble_stripped + "\n\n" + test_code
            if preamble_stripped
            else test_code
        )
        try:
            ast.parse(combined, feature_version=python_version)
            compile(combined, "testgen code", "exec")
            valid_tests.append(test_code)
        except SyntaxError:
            removed_count += 1

    if not valid_tests:
        return preamble_stripped, removed_count

    tests_joined = "\n\n".join(valid_tests)
    result = (
        preamble_stripped + "\n\n\n" + tests_joined
        if preamble_stripped
        else tests_joined
    )
    return result, removed_count


def validate_testgen_code(
    code: str,
    python_version: tuple[int, int],
) -> str:
    """
    Validate generated test code by checking each test individually.
    """
    if not code or code.isspace():
        raise CodeValidationError(
            "No code was generated",
            initial_code=code,
            final_code=code,
            validation_error="Empty code",
        )

    validated_code, tests_removed = _validate_tests_individually(
        code, python_version
    )

    try:
        tree = ast.parse(validated_code, feature_version=python_version)
        compile(validated_code, "testgen code", "exec")

        if not has_test_functions(tree):
            raise CodeValidationError(
                "No test functions found in generated code",
                initial_code=code,
                final_code=validated_code,
                lines_removed=tests_removed,
                validation_error=("No test functions found after validation"),
            )

        return validated_code.strip()

    except SyntaxError as e:
        raise CodeValidationError(
            "Invalid test code: syntax error persists after validation",
            initial_code=code,
            final_code=validated_code,
            lines_removed=tests_removed,
            validation_error=(f"SyntaxError at line {e.lineno}: {e.msg}"),
        ) from e
