# Copyright 2026 Datastrato, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for StdioClientTransport."""

import asyncio
import json
import logging
import sys
import unittest

from adp_sdk.transports.stdio import StdioClientTransport

_ECHO_SCRIPT = "import sys\nfor line in sys.stdin:\n    sys.stdout.write(line)\n    sys.stdout.flush()\n"

_FIXED_RESPONSE_SCRIPT = (
    "import sys, json\n"
    "line = sys.stdin.readline()\n"
    "req = json.loads(line)\n"
    "resp = {'jsonrpc': '2.0', 'id': req['id'], 'result': {'status': 'ok'}}\n"
    "sys.stdout.write(json.dumps(resp) + '\\n')\n"
    "sys.stdout.flush()\n"
)


# ============================================================================


class TestStdioClientTransport(unittest.IsolatedAsyncioTestCase):
    """Tests for StdioClientTransport."""

    async def test_start_and_close(self) -> None:
        """Test basic lifecycle: start and close."""
        transport = StdioClientTransport(sys.executable, ["-c", "import time; time.sleep(10)"])
        await transport.start()
        await transport.write_message('{"test": true}')
        await transport.close()
        with self.assertRaises(RuntimeError):
            await transport.read_message()

    async def test_close_without_start(self) -> None:
        """Close without start should be a no-op."""
        transport = StdioClientTransport(sys.executable, ["-c", "pass"])
        await transport.close()

    async def test_start_twice_raises(self) -> None:
        """Starting an already-started transport should raise."""
        transport = StdioClientTransport(sys.executable, ["-c", "import time; time.sleep(10)"])
        await transport.start()
        try:
            with self.assertRaises(RuntimeError):
                await transport.start()
        finally:
            await transport.close()

    async def test_read_write_echo(self) -> None:
        """Test reading and writing with an echo subprocess."""
        transport = StdioClientTransport(sys.executable, ["-c", _ECHO_SCRIPT])
        async with transport:
            msg = '{"jsonrpc":"2.0","id":1,"method":"test"}'
            await transport.write_message(msg)
            response = await transport.read_message()
            self.assertEqual(response, msg)

    async def test_read_write_json_rpc(self) -> None:
        """Test with a subprocess that returns a JSON-RPC response."""
        transport = StdioClientTransport(sys.executable, ["-c", _FIXED_RESPONSE_SCRIPT])
        async with transport:
            request = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "adp.ping", "params": {}})
            await transport.write_message(request)
            raw = await transport.read_message()
            self.assertIsNotNone(raw)
            data = json.loads(raw)  # type: ignore[arg-type]
            self.assertEqual(data["id"], 1)
            self.assertEqual(data["result"]["status"], "ok")

    async def test_read_message_eof(self) -> None:
        """Test that read_message returns None on EOF."""
        transport = StdioClientTransport(sys.executable, ["-c", "pass"])
        async with transport:
            result = await transport.read_message()
            self.assertIsNone(result)

    async def test_read_before_start_raises(self) -> None:
        """Reading before start should raise."""
        transport = StdioClientTransport(sys.executable, ["-c", "pass"])
        with self.assertRaises(RuntimeError):
            await transport.read_message()

    async def test_write_before_start_raises(self) -> None:
        """Writing before start should raise."""
        transport = StdioClientTransport(sys.executable, ["-c", "pass"])
        with self.assertRaises(RuntimeError):
            await transport.write_message("test")

    async def test_context_manager(self) -> None:
        """Test async context manager usage."""
        transport = StdioClientTransport(sys.executable, ["-c", _ECHO_SCRIPT])
        async with transport as t:
            self.assertIs(t, transport)
            await t.write_message('{"test": true}')
            result = await t.read_message()
            self.assertEqual(result, '{"test": true}')
        with self.assertRaises(RuntimeError):
            await transport.read_message()

    async def test_blank_lines_skipped(self) -> None:
        """Test that blank lines from stdout are skipped."""
        script = "import sys\nsys.stdout.write('\\n\\n')\nsys.stdout.write('{\"ok\":true}\\n')\nsys.stdout.flush()\n"
        transport = StdioClientTransport(sys.executable, ["-c", script])
        async with transport:
            result = await transport.read_message()
            self.assertEqual(result, '{"ok":true}')

    async def test_stderr_forwarded(self) -> None:
        """Test that subprocess stderr is forwarded to logging."""
        script = (
            "import sys\n"
            "sys.stderr.write('test error message\\n')\n"
            "sys.stderr.flush()\n"
            "sys.stdout.write('{\"done\":true}\\n')\n"
            "sys.stdout.flush()\n"
        )
        transport = StdioClientTransport(sys.executable, ["-c", script])
        with self.assertLogs("adp_sdk.transports.stdio.client", level=logging.INFO) as cm:
            async with transport:
                await transport.read_message()
                await asyncio.sleep(0.1)
        self.assertTrue(any("test error message" in msg for msg in cm.output))

    async def test_env_passed_to_subprocess(self) -> None:
        """Test that custom env vars are passed to the subprocess."""
        script = (
            "import os, sys\n"
            "val = os.environ.get('ADP_TEST_VAR', '')\n"
            "sys.stdout.write(val + '\\n')\n"
            "sys.stdout.flush()\n"
        )
        transport = StdioClientTransport(
            sys.executable,
            ["-c", script],
            env={"ADP_TEST_VAR": "hello_adp"},
        )
        async with transport:
            result = await transport.read_message()
            self.assertEqual(result, "hello_adp")

    async def test_multiple_messages(self) -> None:
        """Test sending and receiving multiple messages."""
        transport = StdioClientTransport(sys.executable, ["-c", _ECHO_SCRIPT])
        async with transport:
            for i in range(5):
                msg = json.dumps({"id": i, "data": f"msg-{i}"})
                await transport.write_message(msg)
                response = await transport.read_message()
                self.assertIsNotNone(response)
                data = json.loads(response)  # type: ignore[arg-type]
                self.assertEqual(data["id"], i)

    async def test_subprocess_crash(self) -> None:
        """Test handling of a subprocess that exits unexpectedly."""
        script = "import sys; sys.exit(1)"
        transport = StdioClientTransport(sys.executable, ["-c", script])
        async with transport:
            result = await transport.read_message()
            self.assertIsNone(result)
