import hashlib
import json
import logging
import socket
import tempfile
import unittest
from pathlib import Path
import urllib.parse
from unittest.mock import MagicMock, patch

import yaml
import unideskcprovider as provider

from unideskcprovider import (
    DataChannel,
    ProviderManager,
    ProviderConfig,
    ProviderSession,
    UiBus,
    UpdateManager,
    configure_tcp_keepalive,
    default_data_endpoint,
    encode_data_frame,
    parse_bridge_powershell_command,
    powershell_command_for_log,
    redact_command_for_log,
    update_is_newer,
    windows_path_from_bridge,
)


class ProviderConfigTest(unittest.TestCase):
    def test_loads_short_user_yaml_and_derives_data_port(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.token").write_text("secret\n", encoding="utf-8")
            (root / "provider.yaml").write_text(
                "id: G14-win\n"
                "server: ws://152.53.229.148:18082/ws/provider\n"
                "credentialFile: provider.token\n"
                "pythonRuntime: private\n"
                "autoStart: true\n",
                encoding="utf-8",
            )
            config = ProviderConfig(root / "provider.yaml")
            self.assertEqual(config.id, "G14-win")
            self.assertEqual(config.name, "G14-win")
            self.assertEqual((config.data_host, config.data_port), ("152.53.229.148", 18084))
            self.assertTrue(config.auto_start)
            self.assertEqual(config.python_runtime, "private")
            self.assertEqual(len(config.credential_fingerprint), 12)
            self.assertTrue(config.update_enabled)
            self.assertFalse(config.update_auto_apply)
            self.assertEqual(config.update_interval, 300)
            self.assertEqual(config.update_channel, "stable")
            self.assertEqual(config.update_download_hosts, ("apps.hwpod.com",))
            self.assertEqual(config.update_proxy_mode, "direct")

    def test_rejects_unknown_field(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.yaml").write_text(
                "id: G14-win\nserver: ws://example.test/ws/provider\ncredentialFile: token\nextra: true\n",
                encoding="utf-8",
            )
            with self.assertRaisesRegex(ValueError, "unknown provider.yaml field"):
                ProviderConfig(root / "provider.yaml")

    def test_rejects_invalid_python_runtime_mode(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.token").write_text("secret", encoding="utf-8")
            (root / "provider.yaml").write_text(
                "id: G14-win\n"
                "server: ws://example.test/ws/provider\n"
                "credentialFile: provider.token\n"
                "pythonRuntime: fallback\n",
                encoding="utf-8",
            )
            with self.assertRaisesRegex(ValueError, "pythonRuntime"):
                ProviderConfig(root / "provider.yaml")

    def test_save_auto_start_updates_same_yaml(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.token").write_text("secret", encoding="utf-8")
            path = root / "provider.yaml"
            path.write_text(
                "id: G14-win\nserver: ws://example.test/ws/provider\ncredentialFile: provider.token\n",
                encoding="utf-8",
            )
            config = ProviderConfig(path)
            self.assertEqual(config.python_runtime, "auto")
            config.save_auto_start(True)
            self.assertTrue(yaml.safe_load(path.read_text(encoding="utf-8"))["autoStart"])

    def test_rejects_unknown_update_field(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.token").write_text("secret", encoding="utf-8")
            (root / "provider.yaml").write_text(
                "id: G14-win\n"
                "server: ws://example.test/ws/provider\n"
                "credentialFile: provider.token\n"
                "update:\n"
                "  surprise: true\n",
                encoding="utf-8",
            )
            with self.assertRaisesRegex(ValueError, "unknown provider.yaml update field"):
                ProviderConfig(root / "provider.yaml")

    def test_save_update_settings_keeps_yaml_short(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "provider.token").write_text("secret", encoding="utf-8")
            path = root / "provider.yaml"
            path.write_text(
                "id: G14-win\nserver: ws://example.test/ws/provider\ncredentialFile: provider.token\n",
                encoding="utf-8",
            )
            config = ProviderConfig(path)
            config.save_update_settings(True, True, 600, "stable")
            update = yaml.safe_load(path.read_text(encoding="utf-8"))["update"]
            self.assertEqual(
                update,
                {
                    "enabled": True,
                    "autoApplyWhenIdle": True,
                    "checkIntervalSeconds": 600,
                    "channel": "stable",
                },
            )


class ProtocolHelpersTest(unittest.TestCase):
    def test_windows_provider_declares_32_warm_channels(self):
        self.assertEqual(32, provider.POOL_SIZE)

    def test_linux_uses_native_bash_without_windows_bridge(self):
        with patch.object(provider.sys, "platform", "linux"):
            executable, arguments, shell_name, command_for_log, stdin_script = provider.native_process_command(
                "printf 'linux-native'"
            )
        self.assertEqual(executable, "/bin/bash")
        self.assertEqual(arguments, ["-lc", "printf 'linux-native'"])
        self.assertEqual(shell_name, "bash")
        self.assertEqual(command_for_log, "printf 'linux-native'")
        self.assertFalse(stdin_script)

    def test_linux_labels_disclose_native_headless_runtime(self):
        config = MagicMock()
        manager = ProviderManager(config, MagicMock(), MagicMock())
        with patch.object(provider.sys, "platform", "linux"):
            labels = manager.labels()
        self.assertTrue(labels["linux"])
        self.assertFalse(labels["windows"])
        self.assertEqual(labels["role"], "linux-provider")
        self.assertEqual(labels["attachMode"], "native-linux")
        self.assertEqual(labels["runtime"], "python-headless")

    def test_missing_register_ack_forces_reconnect(self):
        manager = ProviderManager.__new__(ProviderManager)
        manager.stop_event = provider.threading.Event()
        manager._send_heartbeat = MagicMock()
        control = MagicMock()
        control.receive.side_effect = socket.timeout()
        with patch.object(provider, "REGISTER_TIMEOUT_SECONDS", 0):
            with self.assertRaisesRegex(TimeoutError, "register ack timeout"):
                manager._control_messages(control)

    def test_bridge_path_conversion(self):
        self.assertEqual(windows_path_from_bridge("/mnt/c/Windows/System32"), r"C:\Windows\System32")
        self.assertEqual(windows_path_from_bridge(r"D:\Work"), r"D:\Work")

    def test_controlled_powershell_command(self):
        executable, arguments = parse_bridge_powershell_command(
            "'/mnt/c/Windows/System32/WindowsPowerShell/v1.0/powershell.exe' "
            "-NoProfile -EncodedCommand VwByAGkAdABlAC0ATwB1AHQAcAB1AHQAIAAxAA=="
        )
        self.assertEqual(executable, r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe")
        self.assertIn("-EncodedCommand", arguments)

    def test_decodes_actual_powershell_command_for_log(self):
        source = "Get-ChildItem C:\\Users | Select-Object -First 3"
        encoded = provider.base64.b64encode(source.encode("utf-16le")).decode()
        self.assertEqual(
            powershell_command_for_log(["-NoProfile", "-EncodedCommand", encoded]),
            source,
        )

    def test_extracts_inline_powershell_script_from_trans_bootstrap(self):
        source = "$ErrorActionPreference='Stop'; $script = 'Write-Output ''hello'''; Invoke-Expression $script"
        encoded = provider.base64.b64encode(source.encode("utf-16le")).decode()
        self.assertEqual(
            provider.powershell_command_details_for_log(["-EncodedCommand", encoded]),
            ("PowerShell", "Write-Output 'hello'", False),
        )

    def test_extracts_cmd_user_command_from_trans_bootstrap(self):
        source = (
            "$ErrorActionPreference='Stop'; "
            "& 'C:\\Windows\\System32\\cmd.exe' /d /s /c "
            "'chcp 65001>nul && set \"PYTHONUTF8=1\" && "
            "set \"PYTHONIOENCODING=utf-8\" && echo visible-command'; "
            "exit $LASTEXITCODE"
        )
        encoded = provider.base64.b64encode(source.encode("utf-16le")).decode()
        self.assertEqual(
            provider.powershell_command_details_for_log(["-EncodedCommand", encoded]),
            ("cmd", "echo visible-command", False),
        )

    def test_stdin_trans_bootstrap_is_logged_on_eof(self):
        source = "$script = [Console]::In.ReadToEnd(); Invoke-Expression $script"
        encoded = provider.base64.b64encode(source.encode("utf-16le")).decode()
        self.assertEqual(
            provider.powershell_command_details_for_log(["-EncodedCommand", encoded]),
            ("PowerShell", "<stdin-script awaiting payload>", True),
        )
        manager = MagicMock()
        process = MagicMock()
        session = ProviderSession(manager, "stdin-session", MagicMock(), process, True)
        session.input(b"Write-Output 'actual script'")
        session.eof()
        manager.logger.info.assert_called_once()
        self.assertIn("actual script", manager.logger.info.call_args.args[-1])

    def test_redacts_sensitive_values_from_command_log(self):
        source = (
            "$env:PROVIDER_TOKEN='secret-value'; "
            "curl -H 'Authorization: Bearer bearer-value' "
            "--api-key api-value https://example.test"
        )
        redacted = redact_command_for_log(source)
        self.assertNotIn("secret-value", redacted)
        self.assertNotIn("bearer-value", redacted)
        self.assertNotIn("api-value", redacted)
        self.assertGreaterEqual(redacted.count("<redacted>"), 3)

    def test_tray_exit_command_maps_to_exit_action(self):
        self.assertEqual(provider.TrayIcon.action_for_command(provider.TrayIcon.MENU_EXIT), "exit")

    def test_tray_menu_return_value_emits_exit_event(self):
        bus = UiBus()
        tray = provider.TrayIcon(bus, MagicMock())
        tray.hwnd = 1
        user32 = MagicMock()
        user32.TrackPopupMenu.return_value = provider.TrayIcon.MENU_EXIT
        windll = MagicMock()
        windll.user32 = user32
        with patch.object(provider.ctypes, "windll", windll, create=True):
            tray._show_menu()
        self.assertEqual(bus.events.get_nowait(), ("tray", "exit"))
        user32.PostMessageW.assert_called_once_with(1, 0, 0, 0)

    def test_file_log_rotates_daily_with_finite_retention(self):
        with tempfile.TemporaryDirectory() as directory:
            bus = UiBus()
            path = Path(directory) / "provider.log"
            with (
                patch.object(provider, "LOG_DIR", Path(directory)),
                patch.object(provider, "LOG_PATH", path),
            ):
                logger = provider.setup_logger(bus)
            file_handler = next(
                handler
                for handler in logger.handlers
                if isinstance(handler, provider.TimedRotatingFileHandler)
            )
            self.assertEqual(file_handler.when, "MIDNIGHT")
            self.assertEqual(file_handler.interval, 24 * 60 * 60)
            self.assertEqual(file_handler.backupCount, 14)
            for handler in logger.handlers:
                handler.close()

    def test_tray_icon_reflects_state_and_animates_active_wrench(self):
        disconnected = provider.tray_icon_pixels("disconnected", False)
        connected = provider.tray_icon_pixels("connected", False)
        active_first = provider.tray_icon_pixels("connected", True, 0)
        active_next = provider.tray_icon_pixels("connected", True, 1)
        self.assertEqual(len(disconnected), provider.TRAY_ICON_SIZE ** 2 * 4)
        self.assertNotEqual(disconnected, connected)
        self.assertNotEqual(connected, active_first)
        self.assertNotEqual(active_first, active_next)

    def test_data_endpoint_defaults(self):
        parsed = urllib.parse.urlparse("ws://example.test:8081/ws/provider")
        self.assertEqual(default_data_endpoint(parsed), ("example.test", 8082))

    def test_data_frame_has_big_endian_header_length(self):
        frame = encode_data_frame({"type": "data"}, b"abc")
        header_length = int.from_bytes(frame[:4], "big")
        header = frame[4 : 4 + header_length]
        self.assertIn(b'"length":3', header)
        self.assertEqual(frame[-3:], b"abc")

    def test_data_channel_enables_tcp_keepalive(self):
        value = MagicMock()
        configure_tcp_keepalive(value)
        value.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

    def test_data_channel_keeps_writes_blocking_and_polls_reads(self):
        manager = MagicMock()
        manager.config.data_host = "example.test"
        manager.config.data_port = 8082
        manager.config.id = "D601-VM-HOST"
        manager.config.token = "token"
        value = MagicMock()
        value.recv.return_value = b""
        channel = DataChannel(manager, "sshdata-test")

        with (
            patch.object(provider.socket, "create_connection", return_value=value),
            patch.object(provider.select, "select", return_value=([value], [], [])) as select_call,
        ):
            channel._run()

        value.settimeout.assert_called_once_with(None)
        select_call.assert_called_once_with([value], [], [], 1)
        value.sendall.assert_called_once()
        manager.channel_closed.assert_called_once()

    def test_session_exit_waits_for_all_output_pumps_to_drain(self):
        manager = MagicMock()
        process = MagicMock()
        process.wait.return_value = 0
        session = ProviderSession(manager, "session", MagicMock(), process)
        stdout_thread = MagicMock()
        stderr_thread = MagicMock()

        session._wait(stdout_thread, stderr_thread)

        stdout_thread.join.assert_called_once_with()
        stderr_thread.join.assert_called_once_with()
        manager.finish_session.assert_called_once_with("session", 0)

class UpdateManagerTest(unittest.TestCase):
    @staticmethod
    def _source(version: str = "9.9.9") -> bytes:
        return (
            'APP_NAME = "UniDesk Windows Provider"\n'
            f'APP_VERSION = "{version}"\n'
            "class ProviderApp:\n"
            "    pass\n"
        ).encode()

    @staticmethod
    def _manifest(payload: bytes, version: str = "9.9.9") -> dict:
        return {
            "schemaVersion": 1,
            "channel": "stable",
            "latestVersion": version,
            "downloadUrl": "unideskcprovider.py",
            "bytes": len(payload),
            "sha256": hashlib.sha256(payload).hexdigest(),
        }

    def _config(self, root: Path) -> ProviderConfig:
        (root / "provider.token").write_text("secret", encoding="utf-8")
        path = root / "provider.yaml"
        path.write_text(
            "id: G14-win\nserver: ws://example.test/ws/provider\ncredentialFile: provider.token\n",
            encoding="utf-8",
        )
        return ProviderConfig(path)

    def test_version_comparison(self):
        self.assertTrue(update_is_newer("0.2.1", "0.2.0"))
        self.assertFalse(update_is_newer("0.2.0", "0.2"))
        with self.assertRaises(ValueError):
            update_is_newer("v1", "0.2.0")

    def test_manifest_resolves_relative_download(self):
        payload = self._source()
        manifest = UpdateManager.validate_manifest(
            self._manifest(payload),
            "https://apps.hwpod.com/v1/unideskcprovider/update.json",
            "stable",
        )
        self.assertEqual(
            manifest["downloadUrl"],
            "https://apps.hwpod.com/v1/unideskcprovider/unideskcprovider.py",
        )

    def test_published_manifest_matches_provider_script(self):
        root = Path(__file__).resolve().parent
        payload = (root / "unideskcprovider.py").read_bytes()
        raw = json.loads((root / "update.json").read_text(encoding="utf-8"))
        manifest = UpdateManager.validate_manifest(
            raw,
            "https://apps.hwpod.com/v1/unideskcprovider/update.json",
            "stable",
        )
        UpdateManager.validate_payload(payload, manifest)

    def test_payload_rejects_size_hash_syntax_identity_and_version(self):
        payload = self._source()
        manifest = self._manifest(payload)
        UpdateManager.validate_payload(payload, manifest)
        for replacement, pattern in [
            ({**manifest, "bytes": len(payload) + 1}, "byte count"),
            ({**manifest, "sha256": "0" * 64}, "SHA-256"),
        ]:
            with self.assertRaisesRegex(ValueError, pattern):
                UpdateManager.validate_payload(payload, replacement)
        invalid_syntax = b'APP_NAME = "UniDesk Windows Provider"\nclass ProviderApp:\n'
        with self.assertRaisesRegex(ValueError, "Python validation"):
            UpdateManager.validate_payload(invalid_syntax, self._manifest(invalid_syntax))
        wrong_identity = b'APP_VERSION = "9.9.9"\nclass ProviderApp:\n    pass\n'
        with self.assertRaisesRegex(ValueError, "not a UniDesk"):
            UpdateManager.validate_payload(wrong_identity, self._manifest(wrong_identity))
        with self.assertRaisesRegex(ValueError, "version does not match"):
            UpdateManager.validate_payload(payload, self._manifest(payload, "9.9.8"))

    def test_apply_backs_up_replaces_and_records_state(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            config = self._config(root)
            script = root / "unideskcprovider.py"
            original = self._source("0.1.0")
            payload = self._source()
            script.write_bytes(original)
            updater = UpdateManager(
                config,
                logging.getLogger("update-test"),
                UiBus(),
                lambda: True,
                script_path=script,
                update_dir=root / "update",
            )
            updater.available = UpdateManager.validate_manifest(
                self._manifest(payload),
                config.update_metadata_url,
                config.update_channel,
            )
            updater._download = lambda _url: payload
            self.assertEqual(updater.apply(), "9.9.9")
            self.assertEqual(script.read_bytes(), payload)
            self.assertEqual((root / "update" / "unideskcprovider.py.bak").read_bytes(), original)
            state = json.loads((root / "update" / "update-state.json").read_text(encoding="utf-8"))
            self.assertEqual(state["state"], "restarting")
            self.assertEqual(state["targetVersion"], "9.9.9")

    def test_apply_refuses_active_session_and_unlisted_host(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            config = self._config(root)
            script = root / "unideskcprovider.py"
            payload = self._source()
            script.write_bytes(payload)
            updater = UpdateManager(
                config,
                logging.getLogger("update-test"),
                UiBus(),
                lambda: False,
                script_path=script,
                update_dir=root / "update",
            )
            updater.available = UpdateManager.validate_manifest(
                self._manifest(payload),
                config.update_metadata_url,
                config.update_channel,
            )
            with self.assertRaisesRegex(RuntimeError, "active trans"):
                updater.apply()
            updater.is_idle = lambda: True
            updater.available["downloadUrl"] = "https://example.test/unideskcprovider.py"
            with self.assertRaisesRegex(ValueError, "not allowed"):
                updater.apply()

    def test_download_uses_configured_proxy_mode(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            config = self._config(root)
            updater = UpdateManager(
                config,
                logging.getLogger("update-test"),
                UiBus(),
                lambda: True,
                script_path=root / "unideskcprovider.py",
                update_dir=root / "update",
            )
            response = MagicMock()
            response.__enter__.return_value.read.return_value = b"payload"
            opener = MagicMock()
            opener.open.return_value = response
            with patch("unideskcprovider.urllib.request.build_opener", return_value=opener) as build:
                self.assertEqual(updater._download("https://apps.hwpod.com/file.py"), b"payload")
                build.assert_called_once()
                opener.open.assert_called_once()
            config.update_proxy_mode = "system"
            with patch("unideskcprovider.urllib.request.urlopen", return_value=response) as urlopen:
                self.assertEqual(updater._download("https://apps.hwpod.com/file.py"), b"payload")
                urlopen.assert_called_once()


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