"""Concolic test validation and cleanup utilities."""

from __future__ import annotations

import ast
import logging
import os
import re
import subprocess
import uuid
from typing import TYPE_CHECKING

from codeflash_python._compat import SAFE_SYS_EXECUTABLE, codeflash_temp_dir

if TYPE_CHECKING:
    from pathlib import Path

log = logging.getLogger(__name__)

# Known CrossHair limitations that produce invalid Python syntax
# in generated tests:
# - "<locals>" - higher-order functions returning nested functions
# - " object at 0x" - objects with default __repr__
# - "<list_iterator" - iterator objects
CROSSHAIR_KNOWN_LIMITATION_PATTERNS: tuple[str, ...] = (
    "<locals>",
    " object at 0x",
    "<list_iterator",
)


def make_env_with_project_root(
    project_root: str | Path,
) -> dict[str, str]:
    """Return env with *project_root* prepended to PYTHONPATH."""
    env = os.environ.copy()
    root_str = str(project_root)
    pythonpath = env.get("PYTHONPATH", "")
    if pythonpath:
        env["PYTHONPATH"] = f"{root_str}{os.pathsep}{pythonpath}"
    else:
        env["PYTHONPATH"] = root_str
    return env


def is_valid_concolic_test(
    test_code: str,
    project_root: str | None = None,
) -> bool:
    """Validate a concolic test via AST parsing and pytest."""
    try:
        ast.parse(test_code)
    except SyntaxError:
        is_known = any(
            pattern in test_code
            for pattern in CROSSHAIR_KNOWN_LIMITATION_PATTERNS
        )
        if not is_known:
            log.warning(
                "CrossHair generated test with syntax error:\n%s",
                test_code,
            )
        return False

    temp_path = (
        codeflash_temp_dir / f"concolic_test_{uuid.uuid4().hex}.py"
    ).resolve()
    temp_path.write_text(test_code, encoding="utf-8")

    try:
        result = subprocess.run(  # noqa: S603
            [
                SAFE_SYS_EXECUTABLE,
                "-m",
                "pytest",
                "-x",
                "-q",
                temp_path.as_posix(),
            ],
            check=False,
            capture_output=True,
            text=True,
            cwd=project_root,
            timeout=10,
            env=(
                make_env_with_project_root(project_root)
                if project_root
                else None
            ),
        )
    except (subprocess.TimeoutExpired, Exception):  # noqa: BLE001
        return False
    else:
        return result.returncode == 0
    finally:
        temp_path.unlink(missing_ok=True)


class AssertCleanup:
    """Transform assert lines to extract just the function call."""

    def __init__(self) -> None:
        self.assert_re = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
        self.unittest_re = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")

    def transform_asserts(self, code: str) -> str:
        """Replace assert statements with bare function calls."""
        lines = code.splitlines()
        result_lines = []

        for line in lines:
            transformed = self._transform_assert_line(line)
            result_lines.append(
                transformed if transformed is not None else line
            )

        return "\n".join(result_lines)

    def _transform_assert_line(self, line: str) -> str | None:
        indent = line[: len(line) - len(line.lstrip())]

        assert_match = self.assert_re.match(line)
        if assert_match:
            expression = assert_match.group(1).strip()
            if expression.startswith("not "):
                return f"{indent}{expression}"

            expression = expression.rstrip(",;")
            return f"{indent}{expression}"

        unittest_match = self.unittest_re.match(line)
        if unittest_match:
            indent, _assert_method, args = unittest_match.groups()

            if args:
                arg_parts = self._first_top_level_arg(args)
                if arg_parts:
                    return f"{indent}{arg_parts}"

        return None

    def _first_top_level_arg(self, args: str) -> str:
        """Extract the first top-level argument from a call."""
        depth = 0
        for i, ch in enumerate(args):
            if ch in "([{":
                depth += 1
            elif ch in ")]}":
                depth -= 1
            elif ch == "," and depth == 0:
                return args[:i].strip()
        return args.strip()


def clean_concolic_tests(test_suite_code: str) -> str:
    """Clean concolic tests by removing assertions around calls."""
    try:
        tree = ast.parse(test_suite_code)
        can_parse = True
    except Exception:  # noqa: BLE001
        can_parse = False
        tree = None

    if not can_parse or tree is None:
        return AssertCleanup().transform_asserts(test_suite_code)

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
            new_body: list[ast.stmt] = []
            for stmt in node.body:
                if isinstance(stmt, ast.Assert):
                    if isinstance(stmt.test, ast.Compare) and isinstance(
                        stmt.test.left, ast.Call
                    ):
                        new_body.append(ast.Expr(value=stmt.test.left))
                    else:
                        new_body.append(stmt)
                else:
                    new_body.append(stmt)
            node.body = new_body

    return ast.unparse(tree).strip()
