"""Subprocess spawning for test discovery and benchmark execution."""

from __future__ import annotations

import logging
import pickle
import re
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any

from codeflash_python._compat import SAFE_SYS_EXECUTABLE

from ._concolic import make_env_with_project_root

if TYPE_CHECKING:
    from collections.abc import Mapping

log = logging.getLogger(__name__)


def discover_tests_in_subprocess(
    cwd: Path,
    tests_root: Path,
    timeout: int = 300,
) -> tuple[int, list[dict[str, str]], Path | None]:
    """Run pytest test collection in a subprocess and return results.

    Spawns ``_discovery_worker.py`` which runs ``pytest --collect-only``
    and writes collected tests to a temporary pickle file.

    Returns a tuple of ``(exit_code, tests, pytest_rootdir)`` where
    *tests* is a list of dicts with keys ``test_file``, ``test_class``,
    and ``test_function``.  On failure the tuple is ``(-1, [], None)``.
    """
    script_path = Path(__file__).parent / "analysis" / "_discovery_worker.py"

    with tempfile.TemporaryDirectory() as tmpdir:
        pickle_path = str(Path(tmpdir) / "discovery_results.pkl")

        cmd = [
            sys.executable,
            str(script_path),
            str(cwd),
            str(tests_root),
            pickle_path,
        ]
        log.debug("discovering tests with command: %s", " ".join(cmd))

        try:
            subprocess.run(  # noqa: S603
                cmd,
                cwd=str(cwd),
                timeout=timeout,
                check=False,
                text=True,
                capture_output=True,
            )

            pickle_file = Path(pickle_path)
            if not pickle_file.exists():
                log.warning(
                    "discovery pickle file not found: %s",
                    pickle_path,
                )
                return (-1, [], None)

            with pickle_file.open("rb") as f:
                exit_code, tests, rootdir = pickle.load(f)  # noqa: S301

            return (int(exit_code), tests, rootdir)

        except subprocess.TimeoutExpired:
            log.warning(
                "test discovery timed out after %d seconds",
                timeout,
            )
            return (-1, [], None)
        except (
            OSError,
            pickle.UnpicklingError,
            EOFError,
            ValueError,
        ) as exc:
            log.warning("failed to read discovery results: %s", exc)
            return (-1, [], None)


def run_trace_benchmarks_in_subprocess(
    benchmarks_root: Path,
    tests_root: Path,
    trace_file: Path,
    project_root: Path | None = None,
    timeout: int = 600,
) -> subprocess.CompletedProcess[str]:
    """Run benchmark tests with tracing in a subprocess.

    Spawns ``_benchmark_worker.py`` which runs pytest with tracing
    enabled to capture benchmark call traces.  If *project_root* is
    ``None`` the current working directory is used.
    """
    script_path = (
        Path(__file__).parent.parent / "benchmarking" / "_benchmark_worker.py"
    )
    effective_root = project_root or Path.cwd()

    cmd = [
        sys.executable,
        str(script_path),
        str(benchmarks_root),
        str(tests_root),
        str(trace_file),
    ]
    log.debug("running trace benchmarks with command: %s", " ".join(cmd))

    return subprocess.run(  # noqa: S603
        cmd,
        cwd=str(effective_root),
        timeout=timeout,
        check=False,
        text=True,
        capture_output=True,
    )


def get_cross_platform_subprocess_run_args(  # noqa: PLR0913
    *,
    cwd: Path | str | None = None,
    env: Mapping[str, str] | None = None,
    timeout: float | None = None,
    check: bool = False,
    text: bool = True,
    capture_output: bool = True,
) -> dict[str, Any]:
    """Build cross-platform kwargs for ``subprocess.run``."""
    run_args: dict[str, Any] = {
        "cwd": cwd,
        "env": env,
        "text": text,
        "timeout": timeout,
        "check": check,
    }
    if text:
        run_args["errors"] = "replace"
    if sys.platform == "win32":
        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
        run_args["creationflags"] = creationflags
        run_args["stdout"] = subprocess.PIPE
        run_args["stderr"] = subprocess.PIPE
        run_args["stdin"] = subprocess.DEVNULL
    else:
        run_args["capture_output"] = capture_output
    return run_args


def trace_benchmarks_pytest(
    benchmarks_root: Path,
    tests_root: Path,
    project_root: Path,
    trace_file: Path,
    timeout: int = 300,
) -> None:
    """Run benchmark tracing via pytest in a subprocess."""
    benchmark_env = make_env_with_project_root(
        project_root,
    )
    run_args = get_cross_platform_subprocess_run_args(
        cwd=project_root,
        env=benchmark_env,
        timeout=timeout,
        check=False,
        text=True,
        capture_output=True,
    )
    script = (
        Path(__file__).parent.parent / "benchmarking" / "_benchmark_worker.py"
    )
    result = subprocess.run(  # noqa: S603, PLW1510
        [
            SAFE_SYS_EXECUTABLE,
            str(script),
            str(benchmarks_root),
            str(tests_root),
            str(trace_file),
        ],
        **run_args,
    )
    if result.returncode != 0:
        combined_output = result.stdout
        if result.stderr:
            combined_output = (
                combined_output + "\n" + result.stderr
                if combined_output
                else result.stderr
            )

        if "ERROR collecting" in combined_output:
            error_pattern = (
                r"={3,}\s*ERRORS\s*={3,}\n"
                r"([\s\S]*?)(?:={3,}|$)"
            )
            match = re.search(
                error_pattern,
                combined_output,
            )
            error_section = match.group(1) if match else combined_output
        elif "FAILURES" in combined_output:
            error_pattern = (
                r"={3,}\s*FAILURES\s*={3,}\n"
                r"([\s\S]*?)(?:={3,}|$)"
            )
            match = re.search(
                error_pattern,
                combined_output,
            )
            error_section = match.group(1) if match else combined_output
        else:
            error_section = combined_output
        log.warning(
            "Error collecting benchmarks - Pytest Exit code: %s, %s",
            result.returncode,
            error_section,
        )
        log.debug(
            "Full pytest output:\n%s",
            combined_output,
        )
