import os
from pathlib import Path

import pytest

from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.memory_fs import ProjectDir, ProjectFile
from synthegrator.sandboxing import (
    PY_DEFAULT_DOCKER_ENV,
    Cmd,
    pytest_on_docker,
    run_on_docker,
    safer_exec, ExecLimits,
)

IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"


def test_simple_run():
    result = safer_exec("", "3 + 4", user_confirmation=False)
    assert result.output == 7


def test_simple_run_with_code():
    result = safer_exec("foo = lambda x: x*2", "foo(4)", user_confirmation=False)
    assert result.output == 8


def test_with_exit():
    result = safer_exec("exit()", "3 + 4", user_confirmation=False)
    assert result.output is None


def test_with_file_io():
    """Try to write to a temp file. Assert that the file is not written to."""
    import tempfile

    temp_file_name = Path(tempfile.gettempdir()) / "code_for_sandbox_test.txt"
    if temp_file_name.exists():
        temp_file_name.unlink()
    result = safer_exec(
        f"open('{temp_file_name}', 'w').write('hello')",
        "5",
        user_confirmation=False,
    )
    assert result.output is None
    assert result.exception is not None
    assert not temp_file_name.exists()


XML_EXPECTED_EXAMPLE = r"""<?xml version="1.0" encoding="utf-8"?><testsuites name="pytest tests"><testsuite name="pytest" errors="0" failures="1" skipped="0" tests="2" time="0.014" timestamp="2023-07-28T23:17:24.483422" hostname="xxxx"><testcase classname="test_sandboxing" name="test_simple_run" time="0.000" /><testcase classname="test_sandboxing" name="test_simple_run2" time="0.000"><failure message="assert False">def test_simple_run2():
&gt;       assert False
E       assert False

test_sandboxing.py:4: AssertionError</failure></testcase></testsuite></testsuites>"""


def _strip_xml_variance(xml: str) -> str:
    import re

    xml = re.sub(r'timestamp="[^"]+"', "", xml)
    xml = re.sub(r'hostname="[^"]+"', "", xml)
    return re.sub(r'time="[^"]+"', "", xml)


def test_basic_docker_exec():
    r = run_on_docker(
        docker_context=PythonLangSpec().get_default_docker_env(),
        cmds=[
            Cmd("echo hello world"),
        ],
    )
    assert len(r) == 1
    assert r[0].exit_code == 0
    assert r[0].stdout == b"hello world\n"
    assert r[0].completed


def test_basic_docker_exec_with_files():
    r = run_on_docker(
        docker_context=PythonLangSpec().get_default_docker_env(),
        cmds=[
            Cmd("cat foo.txt"),
            Cmd("cat a/bar.txt"),
        ],
        files=ProjectDir(
            {
                "foo.txt": ProjectFile("foo.txt", "hello world"),
                "a": ProjectDir(
                    {
                        "bar.txt": ProjectFile("bar.txt", "bar hello world"),
                        "baz.txt": ProjectFile("baz.txt", "baz hello world"),
                    },
                ),
            },
        ),
    )
    assert len(r) == 2
    assert r[0].exit_code == 0
    assert r[0].stdout == b"hello world"
    assert r[0].completed
    assert r[1].exit_code == 0
    assert r[1].stdout == b"bar hello world"
    assert r[1].completed


def test_pytest_run():
    run_file = "test_sandboxing.py"
    contents = (
        "def test_simple_run():\n"
        "    assert True\n"
        "def test_simple_run2():\n"
        "    assert False\n"
    )
    r = pytest_on_docker(
        files=ProjectDir(
            {
                run_file: ProjectFile(run_file, contents),
            },
        ),
        run_file=run_file,
    )
    assert r.got_results
    assert not r.collection_error
    assert not r.exec_error
    assert r.test_suite_exec_result.exit_code == pytest.ExitCode.TESTS_FAILED
    assert _strip_xml_variance(r.xml_result) == _strip_xml_variance(
        XML_EXPECTED_EXAMPLE,
    )


def test_docker_run_syntax_error():
    run_file = "test_sandboxing.py"
    contents = (
        "foo = 3 +\n"
        "def test_simple_run():\n"
        "    assert True\n"
        "def test_simple_run2():\n"
        "    assert True\n"
    )
    r = pytest_on_docker(
        files=ProjectDir(
            {
                run_file: ProjectFile(run_file, contents),
            },
        ),
        run_file=run_file,
    )
    assert not r.got_results
    assert r.collection_error
    assert not r.exec_error


def test_docker_run_nested_dir():
    r = run_on_docker(
        docker_context=PY_DEFAULT_DOCKER_ENV,
        files=ProjectDir.construct_with_one_file("foo/bar/baz.txt", b"hello world"),
        cmds=[
            "cat foo/bar/baz.txt",
        ],
    )
    assert len(r) == 1
    assert r[0].exit_code == 0
    assert r[0].stdout == b"hello world"


def test_docker_run_timeout():
    """Test that timeouts are properly detected and resources cleaned up."""
    # Create a Python file that contains an infinite loop
    infinite_script = """
import time
while True:
    time.sleep(1)
    print("Still running...")
"""

    files = ProjectDir({
        "infinite_loop.py": ProjectFile("infinite_loop.py", infinite_script)
    })

    # Set a short timeout
    short_limits = ExecLimits(
        timeout_cpu_s=2,
        timeout_realtime_s=2,
        memory_limit_mb=1000,
        max_procs=10  # Allow more processes - important!
    )

    # Run the command
    r = run_on_docker(
        docker_context=PY_DEFAULT_DOCKER_ENV,
        cmds=[Cmd("python infinite_loop.py")],
        files=files,
        limits=short_limits
    )

    # Print diagnostic info
    print(f"Timeout detected: {r[0].timeout}")
    print(f"Completed: {r[0].completed}")
    print(f"Exit code: {r[0].exit_code}")
    print(f"Duration: {r[0].duration}")
    print(f"OOM killed: {r[0].oom_killed}")
    print(f"Stdout: {r[0].stdout}")
    print(f"Stderr: {r[0].stderr}")

    # Assert timeout is correctly detected
    assert r[0].timeout == True
    assert r[0].completed == False