"""End-to-end test utilities for codeflash-python.

Adapted from the original codeflash repo's
``tests/scripts/end_to_end_test_utilities.py``.  Uses attrs and
the codeflash-python CLI entry point instead of the monolith.
"""

from __future__ import annotations

import logging
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import TYPE_CHECKING

import attrs

if TYPE_CHECKING:
    from collections.abc import Callable

log = logging.getLogger(__name__)


@attrs.frozen
class CoverageExpectation:
    """Expected coverage for a single function."""

    function_name: str
    expected_coverage: float = 100.0
    expected_lines: tuple[int, ...] = ()


@attrs.frozen
class E2ETestConfig:
    """Configuration for a single end-to-end optimization test."""

    file_path: Path | None = None
    function_name: str | None = None
    expected_unit_tests_count: int | None = None
    expected_unit_test_files: int | None = None
    min_improvement_x: float = 0.1
    trace_mode: bool = False
    coverage_expectations: tuple[CoverageExpectation, ...] = ()
    benchmarks_root: Path | None = None
    use_worktree: bool = False
    no_gen_tests: bool = False
    expected_acceptance_reason: str | None = None
    expected_in_stdout: tuple[str, ...] = ()


def clear_directory(directory_path: Path) -> None:
    """Remove all files and subdirectories in *directory_path*."""
    if not directory_path.exists():
        log.warning("Directory %s does not exist", directory_path)
        return
    for item in directory_path.iterdir():
        if item.is_file() or item.is_symlink():
            item.unlink()
        elif item.is_dir():
            shutil.rmtree(item)


def cleanup_artifacts(cwd: Path) -> None:
    """Remove pipeline artifacts left behind after an optimization run.

    Cleans up instrumented test files, async wrappers, generated
    test files, temp directories, and other artifacts that the
    optimizer creates in the working tree.
    """
    artifact_globs = [
        "*__perfinstrumented.py",
        "*__perfonlyinstrumented.py",
        "codeflash_async_wrapper.py",
        "*.trace",
        "*.lprof",
        "uv.lock",
    ]
    for pattern in artifact_globs:
        for path in cwd.rglob(pattern):
            path.unlink(missing_ok=True)
            log.debug("Removed artifact: %s", path)

    dir_prefixes = ("tmp", "codeflash_replay_tests_")
    for child in cwd.rglob("*"):
        if child.is_dir() and any(
            child.name.startswith(p) for p in dir_prefixes
        ):
            shutil.rmtree(child, ignore_errors=True)
            log.debug("Removed artifact dir: %s", child)


def build_command(
    cwd: Path,
    config: E2ETestConfig,
    test_root: Path,
) -> list[str]:
    """Build the CLI command for a codeflash-python optimization run."""
    base: list[str] = [
        sys.executable,
        "-m",
        "codeflash_python",
        "--no-pr",
    ]

    if config.file_path is not None:
        base.extend(["--file", str(config.file_path)])
    if config.function_name is not None:
        base.extend(["--function", config.function_name])

    has_config = _has_codeflash_config(cwd)
    if not has_config:
        base.extend(
            [
                "--tests-root",
                str(test_root),
                "--module-root",
                str(cwd),
            ]
        )

    if config.benchmarks_root is not None:
        base.extend(
            [
                "--benchmark",
                "--benchmarks-root",
                str(config.benchmarks_root),
            ]
        )
    if config.use_worktree:
        base.append("--worktree")
    if config.no_gen_tests:
        base.append("--no-gen-tests")

    return base


def _has_codeflash_config(cwd: Path) -> bool:
    """Check whether *cwd* has an existing codeflash config."""
    pyproject = cwd / "pyproject.toml"
    if not pyproject.exists():
        return False

    try:
        import tomllib
    except ModuleNotFoundError:
        import tomli as tomllib  # type: ignore[no-redef]

    try:
        data = tomllib.loads(pyproject.read_text(encoding="utf-8"))
        return "codeflash" in data.get("tool", {})
    except Exception:
        return False


def run_optimization(
    cwd: Path,
    config: E2ETestConfig,
    expected_improvement_pct: int,
) -> bool:
    """Run the full optimization pipeline and validate results.

    Source files are always restored to their original contents
    after the run, regardless of success or failure.  Pipeline
    artifacts (instrumented files, async wrappers, etc.) are
    cleaned up by :func:`cleanup_artifacts`.
    """
    if config.trace_mode:
        return _run_trace_test(
            cwd,
            config,
            expected_improvement_pct,
        )

    path_to_file = cwd / config.file_path if config.file_path else None
    original_contents = (
        path_to_file.read_text(encoding="utf-8")
        if path_to_file and path_to_file.exists()
        else None
    )

    pytest_dir = cwd / "tests" / "pytest"
    test_root = pytest_dir if pytest_dir.is_dir() else cwd / "tests"

    command = build_command(cwd, config, test_root)
    try:
        stdout = _run_subprocess(command, cwd)
        if stdout is None:
            return False

        validated = validate_output(
            stdout,
            expected_improvement_pct,
            config,
        )

        if not validated:
            return False

        if config.expected_in_stdout:
            if not _validate_stdout_contains(
                stdout,
                config.expected_in_stdout,
            ):
                log.error("Expected output not found in candidate")
                return False
            log.info("Expected output found in candidate")

        return validated
    finally:
        if original_contents is not None and path_to_file:
            path_to_file.write_text(
                original_contents,
                encoding="utf-8",
            )
            log.info("Restored %s to original contents", path_to_file)
        cleanup_artifacts(cwd)


def _run_subprocess(
    command: list[str],
    cwd: Path,
) -> str | None:
    """Run *command* and return stdout, or *None* on failure."""
    env = os.environ.copy()
    env["PYTHONIOENCODING"] = "utf-8"

    process = subprocess.Popen(  # noqa: S603
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        cwd=str(cwd),
        env=env,
        encoding="utf-8",
    )

    output: list[str] = []
    assert process.stdout is not None
    for line in process.stdout:
        log.info(line.rstrip())
        output.append(line)

    return_code = process.wait()
    stdout = "".join(output)

    if return_code != 0:
        log.error("Command exited with code %d", return_code)
        return None

    return stdout


def validate_output(  # noqa: PLR0911
    stdout: str,
    expected_improvement_pct: int,
    config: E2ETestConfig,
) -> bool:
    """Validate optimization output meets expected thresholds."""
    if "\u26a1\ufe0f Optimization successful!" not in stdout:
        log.error("Missing optimization success message")
        return False

    improvement_match = re.search(
        r"\U0001f4c8 ([\d,]+)% (?:(\w+) )?improvement",
        stdout,
    )
    if not improvement_match:
        log.error("Could not find improvement percentage")
        return False

    improvement_pct = int(
        improvement_match.group(1).replace(",", ""),
    )
    improvement_x = float(improvement_pct) / 100

    log.info(
        "Performance improvement: %d%% (%.2fx)",
        improvement_pct,
        improvement_x,
    )

    if improvement_pct <= expected_improvement_pct:
        log.error(
            "Improvement %d%% not above %d%%",
            improvement_pct,
            expected_improvement_pct,
        )
        return False

    if improvement_x <= config.min_improvement_x:
        log.error(
            "Improvement %.2fx not above %.2fx",
            improvement_x,
            config.min_improvement_x,
        )
        return False

    if config.expected_acceptance_reason is not None:
        actual = improvement_match.group(2)
        if actual != config.expected_acceptance_reason:
            log.error(
                "Expected reason '%s', got '%s'",
                config.expected_acceptance_reason,
                actual,
            )
            return False

    if not _validate_test_counts(stdout, config):
        return False

    if config.coverage_expectations:
        if not validate_coverage(
            stdout,
            config.coverage_expectations,
        ):
            return False

    if config.no_gen_tests and "Generated 0 tests for" not in stdout:
        log.error("Tests generated even with no_gen_tests flag")
        return False

    return True


def _validate_test_counts(
    stdout: str,
    config: E2ETestConfig,
) -> bool:
    """Validate unit test and file counts in output."""
    if config.expected_unit_tests_count is not None:
        match = re.search(
            r"Discovered (\d+) existing unit tests? "
            r"and \d+ replay tests? in [\d.]+s at",
            stdout,
        )
        if not match:
            log.error("Could not find global unit test count")
            return False
        actual = int(match.group(1))
        if actual != config.expected_unit_tests_count:
            log.error(
                "Expected %d unit tests, found %d",
                config.expected_unit_tests_count,
                actual,
            )
            return False

    if config.expected_unit_test_files is not None:
        match = re.search(
            r"Discovered (\d+) existing unit test files?",
            stdout,
        )
        if not match:
            log.error("Could not find unit test file count")
            return False
        actual = int(match.group(1))
        if actual != config.expected_unit_test_files:
            log.error(
                "Expected %d unit test files, found %d",
                config.expected_unit_test_files,
                actual,
            )
            return False

    return True


def validate_coverage(
    stdout: str,
    expectations: tuple[CoverageExpectation, ...],
) -> bool:
    """Validate coverage data for each expected function."""
    if "CoverageData(" not in stdout:
        log.error("No CoverageData found in output")
        return False

    for expect in expectations:
        pattern = (
            r"(?:main|dependent)_func_coverage="
            r"FunctionCoverage\(name='"
            + re.escape(expect.function_name)
            + r"',\s*coverage=([\d.]+),"
            r"\s*executed_lines=\[(.+?)\],"
        )
        match = re.search(pattern, stdout)
        if not match:
            log.error(
                "No coverage data for %s",
                expect.function_name,
            )
            return False

        coverage = float(match.group(1))
        if coverage != expect.expected_coverage:
            log.error(
                "Coverage %.1f != expected %.1f for %s",
                coverage,
                expect.expected_coverage,
                expect.function_name,
            )
            return False

        executed = tuple(int(x) for x in match.group(2).split(", "))
        if executed != expect.expected_lines:
            log.error(
                "Executed lines %s != expected %s for %s",
                executed,
                expect.expected_lines,
                expect.function_name,
            )
            return False

    return True


def _validate_stdout_contains(
    stdout: str,
    expected: tuple[str, ...],
) -> bool:
    """Check that the best-candidate section contains all *expected* strings."""
    start = stdout.find("INFO     Best candidate")
    end = stdout.find("Best Candidate Explanation")
    candidate_output = stdout[start:end] if start >= 0 else stdout
    return all(exp in candidate_output for exp in expected)


def _run_trace_test(
    cwd: Path,
    config: E2ETestConfig,
    expected_improvement_pct: int,
) -> bool:
    """Run tracer-based E2E test.

    Source files are always restored and artifacts cleaned up
    after the run.
    """
    pytest_dir = cwd / "tests" / "pytest"
    test_root = pytest_dir if pytest_dir.is_dir() else cwd / "tests"

    originals: dict[Path, str] = {}
    for py_file in cwd.glob("*.py"):
        originals[py_file] = py_file.read_text(encoding="utf-8")

    clear_directory(test_root)

    command = [
        "uv",
        "run",
        "-m",
        "codeflash_python",
        "optimize",
        "workload.py",
    ]
    try:
        stdout = _run_subprocess(command, cwd)
        if stdout is None:
            return False

        traced = re.search(
            r"Traced (\d+) function calls successfully",
            stdout,
        )
        if not traced:
            log.error("No traced functions found in output")
            return False
        if int(traced.group(1)) != 8:
            log.error(
                "Expected 8 traced functions, got %s",
                traced.group(1),
            )
            return False

        return validate_output(
            stdout,
            expected_improvement_pct,
            config,
        )
    finally:
        for path, content in originals.items():
            path.write_text(content, encoding="utf-8")
        cleanup_artifacts(cwd)


def run_with_retries(
    test_fn: Callable[..., bool],
    *args: object,
    max_retries: int | None = None,
    retry_delay: int | None = None,
    **kwargs: object,
) -> int:
    """Run *test_fn* with retries, returning an exit code.

    Reads ``MAX_RETRIES`` and ``RETRY_DELAY`` from the environment
    if not passed explicitly.
    """
    if max_retries is None:
        max_retries = int(os.getenv("MAX_RETRIES", "3"))
    if retry_delay is None:
        retry_delay = int(os.getenv("RETRY_DELAY", "5"))

    for attempt in range(1, max_retries + 1):
        log.info("=== Attempt %d of %d ===", attempt, max_retries)

        if test_fn(*args, **kwargs):
            log.info("Test passed on attempt %d", attempt)
            return 0

        log.error("Test failed on attempt %d", attempt)
        if attempt < max_retries:
            log.info("Retrying in %d seconds...", retry_delay)
            time.sleep(retry_delay)

    log.error("Test failed after all retries")
    return 1
