import io
import unittest
from unittest.mock import patch

from rpcscope.progress import ProgressEvent, ProgressReporter


class Clock:
    def __init__(self):
        self.now = 0.0

    def __call__(self):
        return self.now


class Terminal(io.StringIO):
    def isatty(self):
        return True


class ProgressTests(unittest.TestCase):
    def test_redirected_output_has_results_and_no_terminal_codes(self):
        output = io.StringIO()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock)
        reporter(ProgressEvent("start", "ms17-010", "10.0.0.1", total=2))
        clock.now = 1.25
        reporter(ProgressEvent("finish", "ms17-010", "10.0.0.1", "vulnerable", 1, 2))
        reporter(ProgressEvent("summary", "scan", status="complete", completed=2, total=2))
        self.assertIn("[0/2] START ms17-010 10.0.0.1", output.getvalue())
        self.assertIn("[1/2] DONE ms17-010 10.0.0.1 vulnerable (1.2s)", output.getvalue())
        self.assertIn("[2/2] SUMMARY scan complete", output.getvalue())
        self.assertNotIn("\x1b", output.getvalue())
        self.assertNotIn("\r", output.getvalue())

    def test_heartbeat_requires_inactivity_and_an_active_task(self):
        output = io.StringIO()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock, interval=10)
        reporter._render()
        self.assertEqual(output.getvalue(), "")
        reporter(ProgressEvent("start", "discovery", total=3))
        initial = output.getvalue()
        clock.now = 9
        reporter._render()
        self.assertEqual(output.getvalue(), initial)
        clock.now = 10
        reporter._render()
        self.assertIn("[0/3] RUNNING discovery (elapsed 10s)", output.getvalue())
        first_heartbeat = output.getvalue()
        clock.now = 19
        reporter._render()
        self.assertEqual(output.getvalue(), first_heartbeat)
        reporter(ProgressEvent("summary", "scan", status="complete", completed=3, total=3))
        final = output.getvalue()
        clock.now = 99
        reporter._render()
        self.assertEqual(output.getvalue(), final)

    def test_parallel_tasks_and_out_of_order_counts(self):
        output = io.StringIO()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock)
        reporter(ProgressEvent("start", "rpcdump", "10.0.0.1", total=4))
        reporter(ProgressEvent("start", "rpcdump", "10.0.0.2", total=4))
        clock.now = 10
        reporter._render()
        self.assertIn("10.0.0.1 (+1 active)", output.getvalue())
        reporter(ProgressEvent("finish", "rpcdump", "10.0.0.2", "failed", 2, 4))
        reporter(ProgressEvent("finish", "rpcdump", "10.0.0.1", "ok", 1, 4))
        self.assertIn("[2/4] DONE rpcdump 10.0.0.1 ok", output.getvalue())

    def test_silent_discovery_advance_keeps_counts_without_printing_dead_hosts(self):
        output = io.StringIO()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock, interval=10)
        reporter(ProgressEvent("start", "tcp-discovery", total=254))
        initial = output.getvalue()
        clock.now = 9
        reporter(ProgressEvent("advance", "tcp-discovery", "192.0.2.9", "closed", 9, 254))
        self.assertEqual(output.getvalue(), initial)
        clock.now = 10
        reporter._render()
        self.assertIn("[9/254] RUNNING tcp-discovery (elapsed 10s)", output.getvalue())
        self.assertNotIn("192.0.2.9", output.getvalue())
        reporter(ProgressEvent("finish", "tcp-discovery", "192.0.2.10", "alive", 10, 254))
        self.assertIn("[10/254] DONE tcp-discovery 192.0.2.10 alive", output.getvalue())

    def test_tty_advance_does_not_draw_dead_target_or_keep_it_active(self):
        output = Terminal()
        reporter = ProgressReporter(output)
        reporter(ProgressEvent("start", "tcp-discovery", total=254))
        initial = output.getvalue()
        reporter(ProgressEvent("advance", "tcp-discovery", "192.0.2.9", "timeout", 1, 254))
        self.assertEqual(output.getvalue(), initial)
        self.assertNotIn("192.0.2.9", reporter._dashboard(reporter.clock()))
        self.assertIn("[1/254]", reporter._dashboard(reporter.clock()))

    def test_hostless_discovery_advance_preserves_aggregate_heartbeat(self):
        output = io.StringIO()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock, interval=10)
        reporter(ProgressEvent("start", "tcp-discovery", total=254))
        initial = output.getvalue()
        reporter(ProgressEvent("advance", "tcp-discovery", completed=1, total=254))
        self.assertEqual(output.getvalue(), initial)
        clock.now = 10
        reporter._render()
        self.assertIn("[1/254] RUNNING tcp-discovery (elapsed 10s)", output.getvalue())

    def test_untrusted_text_cannot_inject_control_sequences_or_lines(self):
        output = io.StringIO()
        reporter = ProgressReporter(output, mode="plain")
        reporter(ProgressEvent(
            "finish", "check", "10.0.0.1\r\nFAKE", "\x1b[31mvulnerable\x1b[0m",
            detail="one\tline\x00\u202e\nonly",
        ))
        self.assertEqual(len(output.getvalue().splitlines()), 1)
        self.assertIn("10.0.0.1 FAKE vulnerable", output.getvalue())
        self.assertIn("one line only", output.getvalue())
        self.assertNotIn("\x1b", output.getvalue())
        self.assertNotIn("\u202e", output.getvalue())

    def test_tty_dashboard_retains_completed_result(self):
        output = Terminal()
        clock = Clock()
        reporter = ProgressReporter(output, clock=clock, color="never")
        reporter(ProgressEvent("start", "check", "10.0.0.1", total=1))
        self.assertIn("\r\x1b[2K[0/1] RUNNING check 10.0.0.1", output.getvalue())
        reporter(ProgressEvent("finish", "check", "10.0.0.1", "inconclusive", 1, 1))
        self.assertIn("DONE check 10.0.0.1 inconclusive (0.0s)\n", output.getvalue())
        reporter.__exit__(None, None, None)
        self.assertTrue(output.getvalue().endswith("\n"))

    def test_plain_mode_never_uses_ansi_even_on_a_terminal(self):
        output = Terminal()
        reporter = ProgressReporter(output, mode="plain")
        reporter(ProgressEvent("start", "discovery"))
        self.assertNotIn("\x1b", output.getvalue())

    def test_context_starts_and_joins_heartbeat_thread(self):
        output = io.StringIO()
        with patch("rpcscope.progress.threading.Thread") as thread:
            with ProgressReporter(output) as reporter:
                reporter(ProgressEvent("start", "discovery"))
                thread.return_value.start.assert_called_once()
            thread.return_value.join.assert_called_once()
            self.assertTrue(reporter._stop.is_set())
        final = output.getvalue()
        reporter(ProgressEvent("start", "ignored"))
        reporter._render()
        self.assertEqual(output.getvalue(), final)

    def test_quiet_does_not_create_a_thread_or_output(self):
        output = io.StringIO()
        with patch("rpcscope.progress.threading.Thread") as thread:
            with ProgressReporter(output, mode="quiet") as reporter:
                reporter(ProgressEvent("start", "discovery"))
                reporter._render()
            thread.assert_not_called()
        self.assertEqual(output.getvalue(), "")


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