import sys
import threading
import unittest

from damaris.gui.ExperimentHandling import ExperimentHandling
from damaris.gui.ResultHandling import ResultHandling


class _FakeWriter:
    def __init__(self):
        self.sent = []

    def send_next(self, job, quit=False):
        self.sent.append(job)


class _FakeResultIterator:
    def __init__(self):
        self.quit_flag = threading.Event()

    def __iter__(self):
        return iter([])


class TestSandboxStdoutRestore(unittest.TestCase):
    """
    sys.stdout/sys.stderr are process-global: the sandbox redirect in
    run() must be undone even when the script raises or returns early.
    """

    def test_experiment_stdout_restored_after_script_exception(self):
        orig_stdout = sys.stdout
        orig_stderr = sys.stderr
        handled = ExperimentHandling("raise ValueError('boom')", None, None,
                                     log_callback=lambda line: None)
        handled.run()
        self.assertIs(sys.stdout, orig_stdout)
        self.assertIs(sys.stderr, orig_stderr)
        self.assertIsNotNone(handled.raised_exception)

    def test_experiment_stdout_restored_on_success(self):
        orig_stdout = sys.stdout
        writer = _FakeWriter()
        handled = ExperimentHandling("print('hello')\nexperiment = lambda: iter([])",
                                     writer, None, log_callback=lambda line: None)
        handled.run()
        self.assertIs(sys.stdout, orig_stdout)
        self.assertIsNone(handled.raised_exception)

    def test_result_stdout_restored_after_script_exception(self):
        orig_stdout = sys.stdout
        orig_stderr = sys.stderr
        handling = ResultHandling("raise ValueError('boom')", _FakeResultIterator(),
                                  None, log_callback=lambda line: None)
        handling.run()
        self.assertIs(sys.stdout, orig_stdout)
        self.assertIs(sys.stderr, orig_stderr)
        self.assertIsNotNone(handling.raised_exception)

    def test_result_stdout_restored_without_result_function(self):
        orig_stdout = sys.stdout
        handling = ResultHandling("x = 1", _FakeResultIterator(), None,
                                  log_callback=lambda line: None)
        handling.run()
        self.assertIs(sys.stdout, orig_stdout)
        self.assertIsNone(handling.raised_exception)


if __name__ == "__main__":
    unittest.main()