from __future__ import annotations

import asyncio
import os
import sys
import unittest

SRC = os.path.join(os.path.dirname(__file__), "..", "src")
if SRC not in sys.path:
    sys.path.insert(0, SRC)

import remote
from remote import RemoteExecutionError, RemoteHandle, RemoteKernelHandle


class RemoteKernelHandleTests(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        env = {
            **os.environ,
            "PYTHONPATH": SRC + os.pathsep + os.environ.get("PYTHONPATH", ""),
        }
        self.proc = await asyncio.create_subprocess_exec(
            sys.executable,
            "-m",
            "rlm.repl",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        assert self.proc.stdout is not None
        ready_line = await self.proc.stdout.readline()
        self.assertTrue(ready_line.startswith(b'{"event":"ready"'))
        self.handle = RemoteKernelHandle("localhost", self.proc, sys.executable, "/tmp")

    async def asyncTearDown(self):
        await self.handle.close()

    async def test_imports_and_call_lambda(self):
        await self.handle.imports("math")
        res = await self.handle.call(lambda: math.sqrt(144))
        self.assertEqual(res, 12.0)

    async def test_chained_imports(self):
        res = await (await self.handle.imports("platform")).call(lambda: bool(platform.system()))
        self.assertTrue(res)

    async def test_repeated_imports_and_reload(self):
        await self.handle.imports("math")
        self.assertIn("math", self.handle._imported_modules)

        # Second import is a client-side no-op
        await self.handle.imports("math")
        res = await self.handle.call(lambda: math.pi > 3.0)
        self.assertTrue(res)

        # Explicit reload forces re-execution
        await self.handle.imports("math", reload=True)
        res2 = await self.handle.call(lambda: math.pi > 3.0)
        self.assertTrue(res2)

    async def test_call_with_args_and_kwargs(self):
        def multiply_and_add(x, y, offset=0):
            return x * y + offset

        res = await self.handle.call(multiply_and_add, 6, 7, offset=5)
        self.assertEqual(res, 47)

    async def test_call_async_function(self):
        async def async_sum(a, b):
            await asyncio.sleep(0.01)
            return a + b

        res = await self.handle.call(async_sum, 15, 27)
        self.assertEqual(res, 42)

    async def test_call_exception_handling(self):
        with self.assertRaises(RemoteExecutionError) as ctx:
            await self.handle.call(lambda: 1 / 0)
        self.assertIn("ZeroDivisionError", str(ctx.exception))

    async def test_handle_liveness_and_poll(self):
        async def slow_work():
            await asyncio.sleep(0.05)
            return "done!"

        h = self.handle.call(slow_work)
        self.assertIsInstance(h, RemoteHandle)
        self.assertTrue(h.running)
        self.assertIsNone(h.poll())

        res = await h
        self.assertEqual(res, "done!")
        self.assertFalse(h.running)
        self.assertEqual(h.poll(), "done!")

    async def test_handle_output_streaming(self):
        def print_and_return():
            print("step 1: started")
            print("step 2: finished")
            return 123

        h = self.handle.call(print_and_return)
        res = await h
        self.assertEqual(res, 123)
        self.assertIn("step 1: started", h.text())
        self.assertIn("step 2: finished", h.text())
        self.assertIn("step 1: started", h.text())
        self.assertIn("step 2", h.tail(1))

    async def test_first_class_remote_bash_handle(self):
        h = self.handle.bash("echo 'first class remote bash'")
        self.assertIsInstance(h, RemoteHandle)
        res = await h
        self.assertEqual(res.exit_code, 0)
        self.assertIn("first class remote bash", res.text())

    async def test_first_class_remote_bash_handle_with_input(self):
        h = self.handle.bash("cat", input="remote stdin string")
        self.assertIsInstance(h, RemoteHandle)
        res = await h
        self.assertEqual(res.exit_code, 0)
        self.assertEqual(res.text(), "remote stdin string")

    async def test_first_class_remote_bash_handle_with_bash_handle_input(self):
        from rlm.bash import bash as local_bash
        local_h = await local_bash("echo 'piped from local handle'")
        h = self.handle.bash("cat", input=local_h)
        res = await h
        self.assertEqual(res.exit_code, 0)
        self.assertIn("piped from local handle", res.text())

    async def test_async_context_manager(self):
        async with self.handle as rem:
            await rem.imports("sys")
            res = await rem.call(lambda: sys.platform)
            self.assertIsInstance(res, str)
        self.assertTrue(self.handle.closed)


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