from __future__ import annotations

import pickle
import subprocess
import sys
from pathlib import Path
from unittest.mock import patch

from codeflash_python.testing._subprocess_runners import (
    discover_tests_in_subprocess,
    run_trace_benchmarks_in_subprocess,
)


def make_completed_process(
    returncode: int = 0,
    stdout: str = "",
    stderr: str = "",
) -> subprocess.CompletedProcess[str]:
    """Create a CompletedProcess with sensible defaults."""
    return subprocess.CompletedProcess(
        args=["pytest"],
        returncode=returncode,
        stdout=stdout,
        stderr=stderr,
    )


class TestDiscoverTestsInSubprocess:
    """discover_tests_in_subprocess subprocess spawning."""

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_spawns_subprocess_with_correct_command(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Command includes sys.executable, worker script, cwd, tests_root, and pickle path."""
        mock_run.return_value = make_completed_process()

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        discover_tests_in_subprocess(cwd=cwd, tests_root=tests_root)

        cmd = mock_run.call_args[0][0]
        assert sys.executable == cmd[0]
        assert str(cwd) in cmd
        assert str(tests_root) in cmd
        assert len(cmd) >= 4

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_returns_parsed_pickle_results(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Returns the (exit_code, tests, rootdir) tuple from the pickle file."""
        expected_tests = [
            {
                "test_file": "/tests/test_foo.py",
                "test_class": None,
                "test_function": "test_bar",
            },
        ]
        expected_rootdir = Path("/project")
        pickle_data = (0, expected_tests, expected_rootdir)

        def write_pickle_side_effect(cmd, **kwargs):
            """Write pickle data to the pickle path (last arg in cmd)."""
            pickle_path = Path(cmd[-1])
            pickle_path.parent.mkdir(parents=True, exist_ok=True)
            with pickle_path.open("wb") as f:
                pickle.dump(
                    pickle_data,
                    f,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
            return make_completed_process(returncode=0)

        mock_run.side_effect = write_pickle_side_effect

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        exit_code, tests, rootdir = discover_tests_in_subprocess(
            cwd=cwd,
            tests_root=tests_root,
        )

        assert 0 == exit_code
        assert expected_tests == tests
        assert expected_rootdir == rootdir

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_returns_error_on_missing_pickle(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Returns (-1, [], None) when the pickle file is not created."""
        mock_run.return_value = make_completed_process(returncode=0)

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        exit_code, tests, rootdir = discover_tests_in_subprocess(
            cwd=cwd,
            tests_root=tests_root,
        )

        assert -1 == exit_code
        assert [] == tests
        assert rootdir is None

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_returns_error_on_corrupt_pickle(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Returns (-1, [], None) when the pickle file contains corrupt data."""

        def write_corrupt_side_effect(cmd, **kwargs):
            """Write corrupt data to the pickle path."""
            for arg in cmd:
                if ".pkl" in str(arg) or "pickle" in str(arg).lower():
                    pickle_path = Path(arg)
                    pickle_path.parent.mkdir(parents=True, exist_ok=True)
                    pickle_path.write_bytes(b"not valid pickle data")
                    break
            return make_completed_process(returncode=0)

        mock_run.side_effect = write_corrupt_side_effect

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        exit_code, tests, rootdir = discover_tests_in_subprocess(
            cwd=cwd,
            tests_root=tests_root,
        )

        assert -1 == exit_code
        assert [] == tests
        assert rootdir is None

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_default_timeout(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Default timeout is 300 seconds."""
        mock_run.return_value = make_completed_process()

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        discover_tests_in_subprocess(cwd=cwd, tests_root=tests_root)

        call_kwargs = mock_run.call_args
        timeout = call_kwargs.kwargs.get("timeout") or call_kwargs[1].get(
            "timeout"
        )
        assert 300 == timeout

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_custom_timeout(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Custom timeout is passed through to subprocess.run."""
        mock_run.return_value = make_completed_process()

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        discover_tests_in_subprocess(
            cwd=cwd,
            tests_root=tests_root,
            timeout=120,
        )

        call_kwargs = mock_run.call_args
        timeout = call_kwargs.kwargs.get("timeout") or call_kwargs[1].get(
            "timeout"
        )
        assert 120 == timeout

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_cleans_up_pickle_file(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Temp pickle file is cleaned up after reading."""
        pickle_paths_seen: list[Path] = []

        def track_pickle_side_effect(cmd, **kwargs):
            """Write a valid pickle and track the path (last arg in cmd)."""
            pickle_path = Path(cmd[-1])
            pickle_paths_seen.append(pickle_path)
            pickle_path.parent.mkdir(parents=True, exist_ok=True)
            with pickle_path.open("wb") as f:
                pickle.dump(
                    (0, [], None),
                    f,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
            return make_completed_process(returncode=0)

        mock_run.side_effect = track_pickle_side_effect

        cwd = tmp_path / "project"
        cwd.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()

        discover_tests_in_subprocess(cwd=cwd, tests_root=tests_root)

        assert len(pickle_paths_seen) == 1
        assert not pickle_paths_seen[0].exists()


class TestRunTraceBenchmarksInSubprocess:
    """run_trace_benchmarks_in_subprocess subprocess spawning."""

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_spawns_subprocess_with_correct_command(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Command includes sys.executable, worker script, benchmarks_root, tests_root, and trace_file."""
        mock_run.return_value = make_completed_process()

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"

        run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=tmp_path,
        )

        cmd = mock_run.call_args[0][0]
        assert sys.executable == cmd[0]
        assert str(benchmarks_root) in cmd
        assert str(tests_root) in cmd
        assert str(trace_file) in cmd

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_returns_completed_process(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Returns the subprocess.CompletedProcess result."""
        expected = make_completed_process(
            returncode=0,
            stdout="collected 5 items",
        )
        mock_run.return_value = expected

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"

        result = run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=tmp_path,
        )

        assert expected is result
        assert 0 == result.returncode
        assert "collected 5 items" == result.stdout

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_default_timeout(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Default timeout is 600 seconds."""
        mock_run.return_value = make_completed_process()

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"

        run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=tmp_path,
        )

        call_kwargs = mock_run.call_args
        timeout = call_kwargs.kwargs.get("timeout") or call_kwargs[1].get(
            "timeout"
        )
        assert 600 == timeout

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_custom_timeout(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """Custom timeout is passed through to subprocess.run."""
        mock_run.return_value = make_completed_process()

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"

        run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=tmp_path,
            timeout=900,
        )

        call_kwargs = mock_run.call_args
        timeout = call_kwargs.kwargs.get("timeout") or call_kwargs[1].get(
            "timeout"
        )
        assert 900 == timeout

    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_uses_project_root_as_cwd(
        self,
        mock_run,
        tmp_path: Path,
    ) -> None:
        """project_root is passed as cwd to subprocess.run."""
        mock_run.return_value = make_completed_process()

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"
        project_root = tmp_path / "my_project"
        project_root.mkdir()

        run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=project_root,
        )

        call_kwargs = mock_run.call_args
        cwd = call_kwargs.kwargs.get("cwd") or call_kwargs[1].get("cwd")
        assert str(project_root) == cwd

    @patch("codeflash_python.testing._subprocess_runners.Path.cwd")
    @patch("codeflash_python.testing._subprocess_runners.subprocess.run")
    def test_uses_cwd_when_project_root_none(
        self,
        mock_run,
        mock_cwd,
        tmp_path: Path,
    ) -> None:
        """When project_root is None, Path.cwd() is used as cwd."""
        mock_run.return_value = make_completed_process()
        fallback_cwd = tmp_path / "fallback"
        fallback_cwd.mkdir()
        mock_cwd.return_value = fallback_cwd

        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()
        tests_root = tmp_path / "tests"
        tests_root.mkdir()
        trace_file = tmp_path / "trace.json"

        run_trace_benchmarks_in_subprocess(
            benchmarks_root=benchmarks_root,
            tests_root=tests_root,
            trace_file=trace_file,
            project_root=None,
        )

        call_kwargs = mock_run.call_args
        cwd = call_kwargs.kwargs.get("cwd") or call_kwargs[1].get("cwd")
        assert str(fallback_cwd) == cwd
