"""Integration checks using the same client as src/.test/start_client.bat."""

import asyncio
from copy import deepcopy

import pytest


class RunningServer:
    def __init__(self, client_module):
        self.client_module = client_module
        self.clients = []
        self.handlers = []
        self.watchers = []
        self.errors = []

    async def __aenter__(self):
        # Lazy imports keep client-only tests runnable if the server has an error.
        from router import Router
        from session import Sessions
        from transport.tcp.handler import handle

        errors = self.errors
        watchers = self.watchers

        class ObservedSessions(Sessions):
            async def new(self, connection):
                uid = await super().new(connection)
                watchers.append(self[uid].watcher_task)
                return uid

            async def feed(self, uid, packet):
                try:
                    await super().feed(uid, packet)
                except Exception as error:
                    # Preserve the original error before the handler enters
                    # its delayed cleanup and closes the client's transport.
                    errors.append(error)
                    raise

        self.router = Router()
        self.sessions = ObservedSessions(self.router)
        self.router.sessions = self.sessions

        def connected(reader, writer):
            task = asyncio.create_task(handle(reader, writer, self.sessions))
            self.handlers.append(task)
            task.add_done_callback(self.handler_finished)

        self.server = await asyncio.start_server(connected, "127.0.0.1", 0)
        self.port = self.server.sockets[0].getsockname()[1]
        return self

    def handler_finished(self, task):
        if not task.cancelled() and task.exception() is not None:
            self.errors.append(task.exception())

    async def __aexit__(self, exc_type, exc, traceback):
        self.server.close()
        await self.server.wait_closed()
        close_results = await asyncio.gather(
            *(client.close() for client in self.clients), return_exceptions=True,
        )
        # Stop pending recovery-window waits after assertions, without modifying
        # the production timer or waiting 60 seconds for each test connection.
        tasks = self.handlers + self.watchers
        for task in tasks:
            if not task.done():
                task.cancel()
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for result in results:
            if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError):
                self.errors.append(result)
        if exc_type is None and self.errors:
            raise self.errors[0]
        if exc_type is None:
            for result in close_results:
                if isinstance(result, BaseException):
                    raise result

    async def client(
        self, name, inputs=("control",), outputs=("telemetry",),
        node_type="test", **recovery,
    ):
        client = self.client_module.Client(
            name, port=self.port, timeout=1,
            inputs=inputs, outputs=outputs, node_type=node_type,
        )
        self.clients.append(client)
        try:
            await client.connect(**recovery)
        except (TimeoutError, self.client_module.ProtocolError) as error:
            if self.errors:
                raise self.errors[0] from error
            raise
        return client

    async def disconnect(self, client):
        # A half-close proves that the SERVER closes the transport immediately.
        client.writer.write_eof()
        assert await asyncio.wait_for(client.reader.read(), 1) == b""
        await client.close()
        assert self.sessions[client.uid].state == "stop"
        assert client.uid in self.sessions.sessions


def test_server_connection_telemetry_heartbeat_and_disconnect(client_module):
    """Проверяет handshake, метаданные, 10 пакетов, heartbeat и сохранение после EOF."""
    async def scenario():
        async with RunningServer(client_module) as server:
            client = await server.client("rover")
            session = server.sessions[client.uid]
            assert session.state == "air"
            assert session.upf == client.upf
            assert session.static_meta["name"] == "rover"
            assert session.static_meta["type"] == "test"
            assert session.static_meta["in"] == {"control"}
            assert session.static_meta["out"] == {"telemetry"}
            assert session.dynamic_meta["active_in"] == set()
            assert session.dynamic_meta["active_out"] == set()
            assert client.uid in server.router.session_routes

            deadline = session.timeout
            for index in range(10):
                await client.telemetry(temperature=42.0 + index, battery=87 - index)
                assert await client.heartbeat() == {"command": "heartbeat_answer"}
                assert session.last_data == {
                    "telemetry": {"temperature": 42.0 + index, "battery": 87 - index},
                }
                assert session.timeout >= deadline
                deadline = session.timeout
            # Rapid heartbeats may share one monotonic clock tick. Exact renewal
            # is checked separately with a controlled clock.

            await server.disconnect(client)
            assert session.last_data == {"telemetry": {"temperature": 51.0, "battery": 78}}
            assert client.uid in server.router.session_routes

    asyncio.run(asyncio.wait_for(scenario(), 10))


def test_logical_errors_disconnect_before_handshake(client_module):
    """Проверяет закрытие TCP после пяти команд в неправильном состоянии протокола."""
    async def scenario():
        async with RunningServer(client_module) as server:
            client = client_module.Client(port=server.port, timeout=1)
            server.clients.append(client)
            client.reader, client.writer = await asyncio.open_connection("127.0.0.1", server.port)
            # Deliberately answer a handshake that the server has not requested.
            for _ in range(5):
                await client.send("handshake_answer", {"status": True, "state": None})
            assert await asyncio.wait_for(client.reader.read(), 1) == b""
            assert len(server.sessions.sessions) == 1
            session = next(iter(server.sessions.sessions.values()))
            assert session.state == "stop"

    asyncio.run(asyncio.wait_for(scenario(), 5))


def test_recovery_preserves_data_and_moves_both_route_directions(client_module):
    """Проверяет восстановление данных и фактическую передачу по маршрутам в обе стороны."""
    async def scenario():
        async with RunningServer(client_module) as server:
            old = await server.client("old-rover")
            peer = await server.client(
                "peer", inputs=["telemetry"], outputs=["control"], node_type="operator",
            )
            # Route creation has no API yet; explicitly seed both directions.
            server.router.session_routes[old.uid].add(peer.uid)
            server.router.session_routes[peer.uid].add(old.uid)
            await old.telemetry()
            await old.heartbeat()
            assert await peer.recv() == {
                "command": "telemetry", "payload": {"temperature": 42.0, "battery": 87},
            }
            server.sessions[old.uid].dynamic_meta = {
                "active_in": {"control"}, "active_out": {"telemetry"},
            }
            last_data = deepcopy(server.sessions[old.uid].last_data)
            dynamic_meta = deepcopy(server.sessions[old.uid].dynamic_meta)
            await server.disconnect(old)

            restored = await server.client("restored-rover", old_token=old.uid, old_upf=old.upf)
            assert restored.uid != old.uid
            assert server.sessions[restored.uid].state == "air"
            assert server.sessions[restored.uid].last_data == last_data
            assert server.sessions[restored.uid].dynamic_meta == dynamic_meta
            assert old.uid not in server.router.session_routes
            assert set(server.router.session_routes[restored.uid]) == {peer.uid}
            assert set(server.router.session_routes[peer.uid]) == {restored.uid}

            await peer.send("control", {"turn": 1})
            await peer.heartbeat()
            assert await restored.recv() == {"command": "control", "payload": {"turn": 1}}
            await restored.telemetry(43.0, 86)
            await restored.heartbeat()
            assert await peer.recv() == {
                "command": "telemetry", "payload": {"temperature": 43.0, "battery": 86},
            }

    asyncio.run(asyncio.wait_for(scenario(), 10))


@pytest.mark.parametrize("reason", ["active", "wrong_upf", "missing"])
def test_recovery_does_not_take_an_ineligible_session(client_module, reason):
    """Проверяет отказ переноса для активной/отсутствующей сессии и неверного upf."""
    async def scenario():
        async with RunningServer(client_module) as server:
            old = await server.client("old-rover")
            await old.telemetry()
            await old.heartbeat()
            if reason != "active":
                await server.disconnect(old)
            uid = -1 if reason == "missing" else old.uid
            upf = old.upf + 1 if reason == "wrong_upf" else old.upf
            fresh = await server.client("fresh-rover", old_token=uid, old_upf=upf)
            assert server.sessions[fresh.uid].last_data == {}
            assert old.uid in server.router.session_routes
            assert server.sessions[old.uid].state == ("air" if reason == "active" else "stop")

    asyncio.run(asyncio.wait_for(scenario(), 10))


def test_old_recovery_proof_can_only_be_used_once(client_module):
    """Проверяет, что повторный старый UID/upf не забирает уже восстановленные данные."""
    async def scenario():
        async with RunningServer(client_module) as server:
            old = await server.client("old-rover")
            await old.telemetry()
            await old.heartbeat()
            await server.disconnect(old)
            restored = await server.client("restored-rover", old_token=old.uid, old_upf=old.upf)
            expected = {"telemetry": {"temperature": 42.0, "battery": 87}}
            assert server.sessions[restored.uid].last_data == expected
            repeated = await server.client("repeated-rover", old_token=old.uid, old_upf=old.upf)
            assert server.sessions[repeated.uid].last_data == {}
            assert server.sessions[restored.uid].last_data == expected
            assert restored.uid in server.router.session_routes
            assert old.uid not in server.router.session_routes

    asyncio.run(asyncio.wait_for(scenario(), 10))


async def expect_timeout_close(client, session):
    read = asyncio.create_task(client.reader.read())
    try:
        done, _ = await asyncio.wait(
            {read, session.watcher_task}, timeout=1,
            return_when=asyncio.FIRST_COMPLETED,
        )
        assert done, "Воркер не обработал истечение дедлайна"
        if session.watcher_task in done:
            # Expose a crashed watcher rather than just a later socket timeout.
            await session.watcher_task
        assert await asyncio.wait_for(read, 1) == b""
        assert session.state == "stop"
    finally:
        if not read.done():
            read.cancel()
        await asyncio.gather(read, return_exceptions=True)


@pytest.mark.parametrize("state", ["new", "token", "ready", "air"])
@pytest.mark.usefixtures("make_manager")
def test_server_timeout_closes_transport_in_each_phase(client_module, timing, state):
    """Проверяет закрытие настоящего TCP по таймауту на каждом этапе подключения."""
    async def scenario():
        async with timing, RunningServer(client_module) as server:
            client = client_module.Client(port=server.port, timeout=1)
            server.clients.append(client)
            client.reader, client.writer = await asyncio.open_connection("127.0.0.1", server.port)
            async with asyncio.timeout(1):
                while not server.sessions.sessions:
                    await asyncio.sleep(0)
            session = next(iter(server.sessions.sessions.values()))
            if state in {"token", "ready", "air"}:
                await client.send("connection", {"name": "idle-rover", "type": "test", "in": [], "out": ["telemetry"]})
                token = await client.recv()
                assert token["command"] == "token_ask"
            if state in {"ready", "air"}:
                await client.send("token_answer", {"token": None, "upf": None})
                assert await client.recv() == {"command": "handshake_ask"}
            if state == "air":
                await client.send("handshake_answer", {"status": True, "state": None})
                await client.heartbeat()
            assert session.state == state
            request = await timing.sleeping("session", session.watcher_task)
            timing.clock.advance(session.timeout - timing.clock())
            request.release()
            await expect_timeout_close(client, session)
            assert server.sessions[session.uid] is session

    asyncio.run(asyncio.wait_for(scenario(), 5))


@pytest.mark.usefixtures("make_manager")
def test_server_heartbeat_prevents_stale_deadline_disconnect(client_module, timing):
    """Проверяет по TCP продление heartbeat и последующее отключение при тишине."""
    async def scenario():
        async with timing, RunningServer(client_module) as server:
            client = await server.client("rover")
            session = server.sessions[client.uid]
            request = await timing.sleeping("session", session.watcher_task)
            old_deadline = session.timeout
            timing.clock.advance(4)
            await client.heartbeat()
            assert session.timeout > old_deadline
            timing.clock.advance(old_deadline - timing.clock())
            request.release()
            await asyncio.sleep(0)
            if session.watcher_task.done():
                await session.watcher_task
            assert session.state == "air"
            await client.heartbeat()
            request = await timing.sleeping("session", session.watcher_task)
            timing.clock.advance(session.timeout - timing.clock())
            request.release()
            await expect_timeout_close(client, session)
            assert client.uid in server.sessions.sessions

    asyncio.run(asyncio.wait_for(scenario(), 5))