import asyncio
import os
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import AsyncMock, patch, sentinel

import ssh_control


class SSHControlTests(unittest.TestCase):
    def setUp(self):
        self.temp = TemporaryDirectory()
        self.previous = os.environ.get("PRIME_SSH_CONTROL_DIR")
        os.environ["PRIME_SSH_CONTROL_DIR"] = self.temp.name

    def tearDown(self):
        if self.previous is None:
            os.environ.pop("PRIME_SSH_CONTROL_DIR", None)
        else:
            os.environ["PRIME_SSH_CONTROL_DIR"] = self.previous
        self.temp.cleanup()

    def test_control_path_is_stable_short_and_transport_specific(self):
        first = ssh_control.control_path_for("user@example.test")
        second = ssh_control.control_path_for("user@example.test")
        alternate = ssh_control.control_path_for("user@example.test", port=2222)
        self.assertEqual(first, second)
        self.assertNotEqual(first, alternate)
        self.assertLess(len(os.fsencode(first)), 100)
        self.assertEqual(Path(self.temp.name).stat().st_mode & 0o777, 0o700)
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        master_args = ssh_control._master_start_args(connection)
        self.assertIn("ControlMaster=yes", master_args)
        self.assertNotIn("-M", master_args)
        self.assertNotIn("-MNf", master_args)
        self.assertIn("-Nf", master_args)
        self.assertIn("ProxyCommand=false", ssh_control._reuse_only_args())

    def test_ensure_reuses_a_live_master(self):
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        live = ssh_control._Result(("ssh",), 0, "", "", connection.control_path)
        with (
            patch.object(ssh_control, "_control", AsyncMock(return_value=live)),
            patch.object(ssh_control, "_execute", AsyncMock()) as execute,
        ):
            result = asyncio.run(ssh_control._ensure_connection(connection, timeout=10))
        self.assertTrue(result.reused)
        execute.assert_not_awaited()

    def test_ensure_replaces_a_stale_master_before_execution(self):
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        missing = ssh_control._Result(
            ("ssh",), 255, "", "missing", connection.control_path
        )
        live = ssh_control._Result(("ssh",), 0, "", "", connection.control_path)
        started = ssh_control._Result(("ssh",), 0, "", "", connection.control_path)
        with (
            patch.object(
                ssh_control, "_control", AsyncMock(side_effect=[missing, live])
            ) as control,
            patch.object(
                ssh_control, "_execute", AsyncMock(return_value=started)
            ) as execute,
            patch.object(ssh_control, "_remove_stale_socket") as remove,
        ):
            result = asyncio.run(ssh_control._ensure_connection(connection, timeout=10))
        self.assertFalse(result.reused)
        self.assertEqual(control.await_count, 2)
        execute.assert_awaited_once()
        remove.assert_called_once_with(connection.control_path)

    def test_rejects_option_and_destination_injection(self):
        with self.assertRaises(ValueError):
            ssh_control.control_path_for("-oProxyCommand=bad")
        with self.assertRaises(ValueError):
            ssh_control.control_path_for("host", options={"ControlPath": "/tmp/other"})
        with self.assertRaises(ValueError):
            ssh_control.control_path_for("host", options={"ProxyCommand": "other"})
        with self.assertRaises(ValueError):
            ssh_control.run("host name", "true")

    def test_connection_commands_return_native_bash_handle(self):
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        with patch.object(ssh_control, "_bash", return_value=sentinel.handle) as start:
            handle = connection.run("printf '%s\n' hello")
        self.assertIs(handle, sentinel.handle)
        self.assertEqual(handle.command, "ssh user@example.test: printf '%s\n' hello")
        command = start.call_args.args[0]
        self.assertIn("ssh", command)
        self.assertIn(str(connection.control_path), command)
        self.assertIn("user@example.test", command)
        self.assertIn("printf", command)

    def test_run_forwards_stdin_to_the_native_handle(self):
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        with patch.object(ssh_control, "_bash", return_value=sentinel.handle) as start:
            connection.run("cat > out", input="hello")
            ssh_control.run("user@example.test", "cat > out", input=b"\x00 binary")
        self.assertEqual(start.call_args_list[0].kwargs["input"], "hello")
        self.assertEqual(start.call_args_list[1].kwargs["input"], b"\x00 binary")

    def test_copy_and_close_return_native_bash_handles(self):
        connection = ssh_control._connection(
            "user@example.test",
            port=None,
            identity_file=None,
            proxy_jump=None,
            options=None,
            persist="30m",
        )
        with patch.object(
            ssh_control,
            "_bash",
            side_effect=[sentinel.upload, sentinel.download, sentinel.close],
        ) as start:
            upload = connection.copy_to("local file", "/tmp/remote file")
            download = connection.copy_from("/tmp/remote file", "local file")
            close = connection.close()
        self.assertIs(upload, sentinel.upload)
        self.assertIs(download, sentinel.download)
        self.assertIs(close, sentinel.close)
        self.assertEqual(upload.command, "scp local file -> user@example.test:/tmp/remote file")
        self.assertEqual(download.command, "scp user@example.test:/tmp/remote file -> local file")
        self.assertEqual(close.command, "ssh -O exit user@example.test")
        self.assertEqual(start.call_count, 3)
        self.assertIn("scp", start.call_args_list[0].args[0])
        self.assertIn("ssh", start.call_args_list[2].args[0])
        self.assertIn("-O exit", start.call_args_list[2].args[0])


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