"""Integration checks using the same client as examples/client.py."""

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.records = []
        self.errors = []
        self.expected_errors = []
        self.peer_errors = {}

    async def __aenter__(self):
        self.initial_tasks = asyncio.all_tasks()
        # 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
        records = self.records
        peer_errors = self.peer_errors

        class ObservedSessions(Sessions):
            async def new(self, connection):
                uid = await super().new(connection)
                watchers.append(self[uid].watcher_task)
                records.append(self[uid])
                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)
                    peer_errors.setdefault(asyncio.current_task(), []).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())

    def accept_peer_rejection(self, task, error):
        # Accept only the checked exception from this particular client handler.
        assert task in self.handlers and task.done() and not task.cancelled()
        assert task.exception() is error or any(error is item for item in self.peer_errors.get(task, []))
        self.expected_errors.append(error)

    def assert_no_unexpected_errors(self):
        for error in self.errors:
            if not any(error is expected for expected in self.expected_errors):
                raise error

    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 = set(self.handlers + self.watchers)
        tasks.update(asyncio.all_tasks() - self.initial_tasks)
        for task in tasks:
            if not task.done():
                task.cancel()
        # A cancelled handler can enter finally and start the recovery wait.
        # Let it close its transport, then cancel that new wait as well.
        await asyncio.sleep(0)
        for task in tasks:
            if not task.done():
                task.cancel()
        results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 1)
        for result in results:
            if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError):
                self.errors.append(result)
        if exc_type is None:
            self.assert_no_unexpected_errors()
            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))