from __future__ import annotations

import ast
import logging

import isort
import sentry_sdk


def validate_testgen_code(code: str, python_version: tuple[int, int, int], max_lines_to_remove: int = 20) -> str:
    """Validates Python testgen code by iteratively truncating invalid code.

    Args:
        code: String containing Python test code
        python_version: Tuple of (major, minor, micro) Python version
        max_lines_to_remove: Maximum number of lines to remove. If -1, all lines will be considered

    Returns:
        String of valid testgen code

    Raises:
        ValueError: If no valid code is found after checking lines

    """
    fixed_code = isort.code(code, float_to_top=True)
    # Split code into lines and calculate end positions
    lines = fixed_code.splitlines(keepends=True)
    end_positions = []
    current_position = 0
    for line in lines:
        current_position += len(line)
        end_positions.append(current_position)

    truncate_pos = 0
    truncated_code = fixed_code.strip()
    if max_lines_to_remove == -1:
        max_lines_to_remove = len(lines)

    # Iterative remove lines from the end and check for valid code
    for i in range(min(max_lines_to_remove, len(lines)) + 1):
        if not truncated_code or truncated_code.isspace():
            if i == 0:
                raise ValueError("No code was generated when generating tests")
            logging.debug(f"Truncated code is empty after removing {i} lines")
            logging.debug(f"Removed code:\n{fixed_code[truncate_pos:]}")
            raise ValueError(f"Invalid test code generated, failed to validate code even after truncating {i} lines")

        try:
            # Attempt to parse the truncated code
            tree = ast.parse(truncated_code, feature_version=python_version[:2])
            compile(truncated_code, "testgen code", "exec")
            if not has_test_functions(tree):
                logging.debug(f"No test functions found after removing {i} lines")
                logging.debug(f"Removed code:\n{fixed_code[truncate_pos:]}")
                raise ValueError("No test functions were found in the generated code")

            logging.debug(f"Found valid code after removing {i} lines")
            logging.debug(f"Removed code:\n{fixed_code[truncate_pos:]}")
            return truncated_code

        except SyntaxError as e:
            logging.debug(f"Syntax error found: {e}")
            if i >= len(end_positions):
                break
            truncate_pos = end_positions[-(i + 1)]
            truncated_code = fixed_code[:truncate_pos].strip()
            continue
    sentry_sdk.set_tag("error_source", "testgen_code_validation")
    sentry_sdk.set_context(
        "final_error_state",
        {
            "error_type": "validation_failed",
            "total_iterations": min(max_lines_to_remove, len(lines)),
            "initial_testgen_code": code,
            "fixed_testgen-code": fixed_code,
            "final_testgen_code": truncated_code,
        },
    )
    sentry_sdk.capture_exception(
        ValueError(
            f"Invalid test code generated, failed to validate code even after removing {min(max_lines_to_remove, len(lines))} lines"
        )
    )
    raise ValueError(
        f"Invalid test code generated, failed to validate code even after removing {min(max_lines_to_remove, len(lines))} lines"
    )


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

    class TestVisitor(ast.NodeVisitor):
        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            # Check for pytest style (function starting with 'test_')
            if node.name.startswith("test_"):
                raise StopIteration(True)

            # Check for unittest style (class inheriting from TestCase)
            parent: ast.AST | None = getattr(node, "parent", None)
            if isinstance(parent, ast.ClassDef):
                for base in parent.bases:
                    if isinstance(base, ast.Name) and base.id == "TestCase":
                        if node.name.startswith("test_"):
                            raise StopIteration(True)

            self.generic_visit(node)

    visitor = TestVisitor()
    try:
        visitor.visit(tree)
    except StopIteration as found:
        return bool(found.value)
    return False
