from __future__ import annotations

import ast
import logging
import re
from dataclasses import dataclass

import sentry_sdk

from aiservice.common_utils import safe_isort
from core.languages.python.testgen.ast_utils.test_detection import (
    is_test_class,
    is_test_function,
    is_test_function_name,
)


@dataclass
class CodeValidationError(Exception):
    """Exception for code validation failures with debug context."""

    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

    def __post_init__(self) -> None:
        super().__init__(self.message)

    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 validate_testgen_code(code: str, python_version: tuple[int, int]) -> str:
    """Validate Python testgen code by validating each test function individually.

    This function:
    1. Splits the code into preamble (imports, helpers) and test functions
    2. Validates each test individually (preamble + single test)
    3. Keeps only valid tests

    Args:
        code: String containing Python test code
        python_version: Tuple of (major, minor) Python version

    Returns:
        String of valid testgen code

    Raises:
        CodeValidationError: If no valid test code is found

    """
    fixed_code = safe_isort(code, float_to_top=True)

    if not fixed_code or fixed_code.isspace():
        raise CodeValidationError(
            message="No code was generated when generating tests",
            initial_code=code,
            fixed_code=fixed_code,
            final_code=fixed_code,
            lines_removed=0,
            validation_error="Empty code",
        )

    # Try to validate each test function individually
    validated_code, tests_removed = validate_tests_individually(fixed_code, python_version)

    # Check if the validated code is parseable and has test functions
    try:
        tree = ast.parse(validated_code, feature_version=python_version)
        compile(validated_code, "testgen code", "exec")

        if not has_test_functions(tree):
            logging.warning(
                "No test functions found after validation. tests_removed=%d, validated_code_len=%d",
                tests_removed,
                len(validated_code),
            )
            logging.warning("Validated code (first 2000 chars):\n%s", validated_code[:2000])
            raise CodeValidationError(
                message="No test functions were found in the generated code",
                initial_code=code,
                fixed_code=fixed_code,
                final_code=validated_code,
                lines_removed=tests_removed,
                validation_error="No test functions found after individual validation",
            )

        if tests_removed > 0:
            logging.info("Removed %d invalid tests during individual validation", tests_removed)

        return validated_code.strip()

    except SyntaxError as e:
        # If we still have syntax errors after individual validation,
        # it means the preamble itself has errors or splitting failed
        last_syntax_error = f"SyntaxError at line {e.lineno}: {e.msg}"
        logging.warning("Syntax error in validated code: %s", e)

        sentry_sdk.set_tag("error_source", "testgen_code_validation")
        sentry_sdk.set_context(
            "final_error_state",
            {
                "error_type": "validation_failed",
                "tests_removed": tests_removed,
                "initial_testgen_code": code,
                "fixed_testgen_code": fixed_code,
                "final_testgen_code": validated_code,
            },
        )

        err = CodeValidationError(
            message="Invalid test code generated, syntax error persists after validation",
            initial_code=code,
            fixed_code=fixed_code,
            final_code=validated_code,
            lines_removed=tests_removed,
            validation_error=last_syntax_error,
        )
        sentry_sdk.capture_exception(err)
        raise err from e


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

    class TestVisitor(ast.NodeVisitor):
        def _check_test_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
            """Check if a function (sync or async) is a test function."""
            if is_test_function_name(node.name):
                raise StopIteration(True)  # noqa: FBT003

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self._check_test_function(node)
            self.generic_visit(node)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            self._check_test_function(node)
            self.generic_visit(node)

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


@dataclass
class CodeParts:
    """Parsed parts of test code."""

    preamble: str  # imports, constants, helper functions/classes
    test_functions: list[str]  # individual test function source code


def split_code_with_ast(code: str, python_version: tuple[int, int]) -> CodeParts | None:
    """Try to split code using AST parsing. Returns None if parsing fails."""
    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):
            # Include decorators - use first decorator's line if present
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.decorator_list:
                start_line = node.decorator_list[0].lineno - 1  # AST uses 1-based line numbers
            else:
                start_line = node.lineno - 1
            end_line = node.end_lineno if node.end_lineno else start_line + 1

            # Capture any gap as part of preamble (comments, blank lines before test)
            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)

            test_code = "".join(lines[start_line:end_line])
            test_functions.append(test_code)
            last_preamble_end = end_line
        else:
            # Include decorators for non-test functions too (helper classes, etc.)
            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 CodeParts(preamble=preamble, test_functions=test_functions)


# Regex patterns for finding test function boundaries (including indented class methods)
# Use [ \t]* instead of \s* to avoid matching newlines
_TEST_FUNCTION_PATTERN = re.compile(
    r"^[ \t]*((?:async\s+)?def\s+test_\w+\s*\([^)]*\)\s*(?:->\s*[^:]+)?:)", re.MULTILINE
)

# Pattern to detect test classes:
# - Classes starting with "Test" (e.g., TestFoo, TestBasic)
# - Classes inheriting from TestCase or unittest.TestCase (e.g., class MyTests(TestCase):)
_TEST_CLASS_PATTERN = re.compile(
    r"^(class\s+(?:Test\w*)\s*(?:\([^)]*\))?\s*:|class\s+\w+\s*\(\s*(?:\w+\.)?TestCase\s*\)\s*:)", re.MULTILINE
)

# Pattern to detect any definition at same or lower indent level (ends the previous function/class)
_ANY_DEF_PATTERN = re.compile(r"^\s*((?:async\s+)?def\s+|class\s+)", re.MULTILINE)


def _get_indent_level(line: str) -> int:
    """Get the indentation level (number of leading spaces/tabs) of a line."""
    return len(line) - len(line.lstrip())


def find_decorator_start_line(lines: list[str], def_line: int) -> int:
    """Find the first line of decorators preceding a function/class definition.

    Looks backward from def_line to find consecutive decorator lines,
    skipping blank lines and comments between decorators.

    Args:
        lines: List of source code lines
        def_line: 0-based line index of the function/class definition

    Returns:
        The line index where the decorator block starts, or def_line if no decorators found.

    """
    if def_line <= 0:
        return def_line

    start_line = def_line
    i = def_line - 1
    max_look_back = 20  # Safety limit to prevent scanning too far

    while i >= 0 and (def_line - i) <= max_look_back:
        line = lines[i].strip()
        if line.startswith("@"):
            # Found a decorator, update start line
            start_line = i
            i -= 1
        elif not line or line.startswith("#"):
            # Skip blank lines and comments between decorators
            i -= 1
        else:
            # Hit non-decorator, non-blank, non-comment line
            break

    return start_line


def split_code_with_regex(code: str) -> CodeParts:
    """Split code using regex when AST parsing fails.

    This is a fallback that finds test function/class boundaries by pattern matching.
    Less precise than AST but works on code with syntax errors.

    Handles both:
    - Top-level test functions (def test_*)
    - Test classes (class Test* or class *TestCase)
    """
    lines = code.splitlines(keepends=True)

    # Find test classes (these are extracted as whole units)
    class_matches = list(_TEST_CLASS_PATTERN.finditer(code))

    # Find top-level test functions (indent level 0)
    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 CodeParts(preamble=code, test_functions=[])

    # Combine all test items (classes and top-level functions) sorted by position
    test_items: list[tuple[int, int, str]] = []  # (start_pos, start_line, type)

    for m in class_matches:
        # Find start of line to ensure we count complete lines
        line_start = code.rfind("\n", 0, m.start()) + 1
        def_line = code[:line_start].count("\n")
        # Look backward for decorators
        start_line = find_decorator_start_line(lines, def_line)
        test_items.append((m.start(), start_line, "class"))

    for m in top_level_test_funcs:
        # Find start of line to ensure we count complete lines
        line_start = code.rfind("\n", 0, m.start()) + 1
        def_line = code[:line_start].count("\n")
        # Look backward for decorators
        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])

    # Find boundaries for each test item
    test_boundaries: list[tuple[int, int]] = []

    for i, (_start_pos, start_line, _item_type) in enumerate(test_items):
        # End at the next test item or end of file
        end_line = test_items[i + 1][1] if i + 1 < len(test_items) else len(lines)
        test_boundaries.append((start_line, end_line))

    # Split into preamble and tests
    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()

        # Trim trailing blank lines from this test
        test_lines = test_code.splitlines(keepends=True)
        while test_lines and not test_lines[-1].strip():
            test_lines.pop()

        # Trim trailing incomplete definitions (e.g., "def" or "class" without body)
        while test_lines:
            last_line = test_lines[-1].rstrip()
            # Check for incomplete definition patterns
            if re.match(r"^\s*(def|class|async\s+def)\s*$", last_line):
                test_lines.pop()
            elif re.match(r"^\s*(def|class|async\s+def)\s+\w*$", last_line):
                # "def foo" or "class Bar" without colon
                test_lines.pop()
            else:
                break

        if test_lines:
            test_functions.append("".join(test_lines).rstrip())

    return CodeParts(preamble=preamble, test_functions=test_functions)


def split_code_into_parts(code: str, python_version: tuple[int, int]) -> CodeParts:
    """Split code into preamble (imports, constants, helpers) and individual test functions.

    First tries AST-based splitting for accuracy, falls back to regex-based
    splitting if the code has syntax errors.

    Args:
        code: The complete test code
        python_version: Tuple of (major, minor) Python version

    Returns:
        CodeParts with preamble and list of test function strings

    """
    # Try AST-based splitting first (more accurate)
    result = split_code_with_ast(code, python_version)
    if result is not None:
        return result

    # Fall back to regex-based splitting for code with syntax errors
    logging.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 by removing offending lines.

    Iteratively removes lines that cause syntax errors. Handles multi-line
    constructs (unclosed parens/brackets/strings) by removing contiguous
    blocks of broken lines.

    Args:
        preamble: The preamble code with potential syntax errors
        python_version: Tuple of (major, minor) Python version

    Returns:
        The repaired preamble (may be empty if nothing is salvageable)

    """
    lines = preamble.splitlines(keepends=True)
    max_attempts = min(len(lines), 15)  # Safety limit

    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 ""
            # Remove the offending line (1-based lineno -> 0-based index)
            error_idx = e.lineno - 1
            if 0 <= error_idx < len(lines):
                lines.pop(error_idx)
            else:
                return ""

    # If we exhausted attempts, return whatever is left
    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 and return only valid tests.

    Args:
        code: The complete test code
        python_version: Tuple of (major, minor) Python version

    Returns:
        Tuple of (valid_code, number_of_tests_removed)

    """
    parts = split_code_into_parts(code, python_version)

    if not parts.test_functions:
        # No test functions found, return original code
        logging.debug("split_code_into_parts found no test functions")
        return code, 0

    logging.debug("split_code_into_parts found %d test functions", len(parts.test_functions))

    valid_tests: list[str] = []
    removed_count = 0

    preamble_stripped = parts.preamble.strip()

    # Check if preamble itself is valid; if not, try to repair it
    if preamble_stripped:
        try:
            ast.parse(preamble_stripped, feature_version=python_version)
        except SyntaxError as e:
            logging.warning("Preamble has syntax error at line %s: %s", e.lineno, e.msg)
            logging.debug("Preamble content:\n%s", preamble_stripped[:2000])
            repaired = repair_preamble(preamble_stripped, python_version)
            if repaired != preamble_stripped:
                logging.info(
                    "Repaired preamble: removed %d chars of broken code", len(preamble_stripped) - len(repaired)
                )
                preamble_stripped = repaired

    for i, test_code in enumerate(parts.test_functions):
        # Combine preamble with this single test
        combined = preamble_stripped + "\n\n" + test_code if preamble_stripped else test_code

        try:
            # Try to parse and compile the combined code
            ast.parse(combined, feature_version=python_version)
            compile(combined, "testgen code", "exec")
            # If successful, keep this test
            valid_tests.append(test_code)
        except SyntaxError as e:
            # This test is invalid, skip it
            removed_count += 1
            # Log the first few failures at WARNING level to help debug
            if removed_count <= 3:
                first_line = test_code.split("\n")[0] if test_code else "(empty)"
                logging.debug(
                    "Removing test %d/%d due to syntax error at line %s: %s. Test starts with: %s",
                    i + 1,
                    len(parts.test_functions),
                    e.lineno,
                    e.msg,
                    first_line[:80],
                )

    if not valid_tests:
        # All tests were invalid
        return preamble_stripped, removed_count

    # Combine preamble with all valid tests
    # Handle empty preamble case (just "\n" or whitespace)
    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
