import asyncio
import os
import tempfile
import threading
import unittest
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from unittest.mock import Mock
from uuid import uuid4

from websockets.asyncio.server import serve

from archive_clients.config import ClientConfig, ConnectionConfig, ServiceConfig
from archive_clients.daemon import ArchiveClientDaemon
from archive_clients.probes import FilesystemProbe
from archive_clients.services import ServiceProbe
from archive_clients.syncthing import ConfiguredRoute
from archive_clients.protocol import decode, encode, encode_message, new_envelope
from archive_control.v1 import (
    client_pb2, common_pb2, control_pb2, inventory_pb2, job_pb2, route_pb2,
)


class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
    async def test_unrelated_job_commands_execute_concurrently(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            service = ServiceConfig("http://local", PurePosixPath("/api"), root)
            config = ClientConfig(
                "cache-1", "Cache 1", "cache", "ws://control", token,
                root / "state.db", root / "backups", service, service,
            )
            probe = FilesystemProbe(root, True, True, True, True, True)
            daemon = ArchiveClientDaemon(config, [probe, probe], [])
            started: set[str] = set()
            release = asyncio.Event()

            async def execute(command, correlation_id, outbound):
                started.add(command.assign_job.job.job_id)
                await release.wait()

            daemon._execute_job_command_locked = execute
            commands = []
            for _ in range(2):
                command = control_pb2.Command(command_id=str(uuid4()))
                command.assign_job.job.job_id = str(uuid4())
                commands.append(command)
            outbound = asyncio.Queue()
            tasks = [
                asyncio.create_task(
                    daemon._execute_job_command(command, "", outbound)
                )
                for command in commands
            ]
            for _ in range(100):
                if len(started) == 2:
                    break
                await asyncio.sleep(0.01)
            self.assertEqual(len(started), 2)
            release.set()
            await asyncio.gather(*tasks)

    async def test_silent_control_connection_ends_for_reconnect(self):
        """A lost server heartbeat must not leave durable commands stranded."""

        async def control(websocket):
            registration = decode(await websocket.recv())
            self.assertEqual(registration.WhichOneof("payload"), "register_request")
            response = new_envelope()
            response.correlation_id = registration.message_id
            response.register_response.status = (
                client_pb2.REGISTRATION_STATUS_ACCEPTED
            )
            response.register_response.negotiated_version.major = 1
            await websocket.send(encode(response))
            # Deliberately keep TCP/WebSocket open but send no application
            # heartbeats. This models a stale proxy/server-side session.
            await websocket.wait_closed()

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
                port = server.sockets[0].getsockname()[1]
                service = ServiceConfig(
                    "http://local", PurePosixPath("/api"), root,
                )
                config = ClientConfig(
                    "cache-1", "Cache 1", "cache",
                    f"ws://127.0.0.1:{port}", token,
                    root / "state.db", root / "backups", service, service,
                    ConnectionConfig(
                        registration_timeout=1,
                        offline_timeout=0.05,
                        reconnect_initial=0.01,
                        reconnect_max=0.01,
                        reconnect_jitter=False,
                    ),
                )
                probe = FilesystemProbe(root, True, True, True, True, True)
                daemon = ArchiveClientDaemon(config, [probe, probe], [])
                await asyncio.to_thread(daemon.store.initialize)
                with self.assertRaisesRegex(
                    RuntimeError, "control heartbeat timed out"
                ):
                    await asyncio.wait_for(daemon._connection(), 1)

    async def test_failed_writer_ends_connection_without_waiting_for_heartbeat(self):
        """A send failure must immediately reach the reconnect supervisor."""

        async def control(websocket):
            registration = decode(await websocket.recv())
            response = new_envelope()
            response.correlation_id = registration.message_id
            response.register_response.status = (
                client_pb2.REGISTRATION_STATUS_ACCEPTED
            )
            response.register_response.negotiated_version.major = 1
            await websocket.send(encode(response))
            await websocket.wait_closed()

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
                port = server.sockets[0].getsockname()[1]
                service = ServiceConfig("http://local", PurePosixPath("/api"), root)
                config = ClientConfig(
                    "cache-1", "Cache 1", "cache",
                    f"ws://127.0.0.1:{port}", token,
                    root / "state.db", root / "backups", service, service,
                )
                probe = FilesystemProbe(root, True, True, True, True, True)
                daemon = ArchiveClientDaemon(config, [probe, probe], [])
                await asyncio.to_thread(daemon.store.initialize)

                async def failed_writer(websocket, outbound):
                    raise OSError("simulated broken socket")

                daemon._writer = failed_writer
                with self.assertRaisesRegex(RuntimeError, "writer failed"):
                    await asyncio.wait_for(daemon._connection(), 1)

    async def test_full_outbound_queue_aborts_connection_instead_of_blocking_heartbeats(self):
        """Bulk output cannot indefinitely block the receive/heartbeat loop."""

        async def control(websocket):
            registration = decode(await websocket.recv())
            response = new_envelope()
            response.correlation_id = registration.message_id
            response.register_response.status = (
                client_pb2.REGISTRATION_STATUS_ACCEPTED
            )
            response.register_response.negotiated_version.major = 1
            await websocket.send(encode(response))
            for sequence in range(1, 102):
                heartbeat = new_envelope()
                heartbeat.heartbeat.sequence = sequence
                await websocket.send(encode(heartbeat))
            await websocket.wait_closed()

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
                port = server.sockets[0].getsockname()[1]
                service = ServiceConfig("http://local", PurePosixPath("/api"), root)
                config = ClientConfig(
                    "cache-1", "Cache 1", "cache",
                    f"ws://127.0.0.1:{port}", token,
                    root / "state.db", root / "backups", service, service,
                    ConnectionConfig(outbound_enqueue_timeout=0.01),
                )
                probe = FilesystemProbe(root, True, True, True, True, True)
                daemon = ArchiveClientDaemon(config, [probe, probe], [])
                await asyncio.to_thread(daemon.store.initialize)

                async def stopped_writer(websocket, outbound):
                    await asyncio.Event().wait()

                daemon._writer = stopped_writer
                with self.assertRaisesRegex(RuntimeError, "outbound queue is blocked"):
                    await asyncio.wait_for(daemon._connection(), 2)

    async def test_reconnect_resume_skips_historical_job_commands(self):
        """Registration reconciliation, not command replay, recovers job state."""

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            service = ServiceConfig("http://local", PurePosixPath("/api"), root)
            config = ClientConfig(
                "cache-1", "Cache 1", "cache", "ws://control", token,
                root / "state.db", root / "backups", service, service,
            )
            probe = FilesystemProbe(root, True, True, True, True, True)
            daemon = ArchiveClientDaemon(config, [probe, probe], [])
            await asyncio.to_thread(daemon.store.initialize)
            command = control_pb2.Command(command_id=str(uuid4()))
            command.execute_step.job_id = str(uuid4())
            command.execute_step.expected_last_event_sequence = 1
            acknowledgement = control_pb2.CommandAck(
                command_id=command.command_id,
                status=control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
            )
            await asyncio.to_thread(
                daemon.store.accept_command,
                command.command_id,
                encode_message(command),
                encode_message(acknowledgement),
            )
            daemon.jobs = Mock()
            outbound = asyncio.Queue()
            tasks = set()
            await daemon._resume_commands(outbound, tasks)
            self.assertEqual(tasks, set())
            self.assertTrue(outbound.empty())

    async def test_eviction_assignment_and_steps_are_admitted(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            service = ServiceConfig(
                "http://local", PurePosixPath("/api"), root
            )
            config = ClientConfig(
                "cache-1", "Cache 1", "cache", "ws://control", token,
                root / "state.db", root / "backups", service, service,
            )
            probe = FilesystemProbe(root, True, True, True, True, True)
            daemon = ArchiveClientDaemon(config, [probe, probe], [])
            daemon.jobs = Mock()

            assign = control_pb2.Command(command_id=str(uuid4()))
            definition = assign.assign_job.job
            definition.job_id = str(uuid4())
            definition.eviction.cache_client_id = "cache-1"
            acknowledgement = daemon._initial_acknowledgement(
                assign, set(), False
            )
            self.assertEqual(
                acknowledgement.status,
                control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
            )

            for step_kind in (
                job_pb2.JOB_STEP_KIND_VERIFY_ARCHIVE_COVERAGE,
                job_pb2.JOB_STEP_KIND_QB_REMOVE_ENTRY,
                job_pb2.JOB_STEP_KIND_SAFE_FILE_UNLINK,
            ):
                execute = control_pb2.Command(command_id=str(uuid4()))
                execute.execute_step.job_id = definition.job_id
                execute.execute_step.step = step_kind
                acknowledgement = daemon._initial_acknowledgement(
                    execute, set(), False
                )
                self.assertEqual(
                    acknowledgement.status,
                    control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
                )

            cancel = control_pb2.Command(command_id=str(uuid4()))
            cancel.cancel_job.job_id = definition.job_id
            acknowledgement = daemon._initial_acknowledgement(
                cancel, set(), False
            )
            self.assertEqual(
                acknowledgement.status,
                control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
            )

    async def test_ensure_route_is_durable_and_duplicate_replays_updates(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            service = ServiceConfig(
                "http://local", PurePosixPath("/sync"), root
            )
            config = ClientConfig(
                "cache-1", "Cache 1", "cache", "ws://control", token,
                root / "state.db", root / "backups", service, service,
            )
            local_route = route_pb2.LocalRoute(
                route_id="route-1",
                local_relative_path="routes/route-1",
                folder_type=route_pb2.SYNCTHING_FOLDER_TYPE_SEND_RECEIVE,
                local_syncthing_device_id="LOCAL",
                peer_syncthing_device_ids=["PEER"],
                state=route_pb2.ROUTE_STATE_PROVISIONING,
                writable=True,
            )
            local_route.observed_at.GetCurrentTime()
            manager = Mock()
            manager.configure.return_value = ConfiguredRoute(
                "LOCAL", root / "routes/route-1", local_route, True, True
            )
            manager.verify.return_value = (True, True)
            probe = FilesystemProbe(root, True, True, True, True, True)
            syncthing_probe = ServiceProbe(
                "syncthing",
                common_pb2.HEALTH_STATE_HEALTHY,
                datetime.now(timezone.utc),
                device_id="LOCAL",
            )
            daemon = ArchiveClientDaemon(
                config,
                [probe, probe],
                [syncthing_probe],
                route_manager=manager,
            )
            await asyncio.to_thread(daemon.store.initialize)
            command = new_envelope()
            command.command.command_id = str(uuid4())
            command.command.created_at.CopyFrom(command.sent_at)
            spec = command.command.ensure_route.route
            spec.route_id = "route-1"
            spec.peer_client_id = "archive-1"
            spec.peer_syncthing_device_id = "PEER"
            spec.peer_addresses.append("dynamic")
            spec.local_relative_path = "routes/route-1"
            spec.setup_timeout_seconds = 1800
            outbound = asyncio.Queue()
            tasks = set()
            await daemon._handle(command, outbound, tasks)
            self.assertEqual(
                decode(await outbound.get()).command_ack.status,
                control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
            )
            await next(iter(tasks))
            first_updates = [decode(await outbound.get()).route_update for _ in range(3)]
            self.assertEqual(
                [update.sequence for update in first_updates], [1, 2, 3]
            )
            self.assertEqual(
                first_updates[-1].verification.state,
                route_pb2.ROUTE_STATE_READY,
            )
            self.assertTrue(first_updates[-1].route.archive_control_created)

            duplicate = new_envelope()
            duplicate.command.CopyFrom(command.command)
            tasks = set()
            await daemon._handle(duplicate, outbound, tasks)
            self.assertEqual(
                decode(await outbound.get()).command_ack.status,
                control_pb2.COMMAND_ACK_STATUS_DUPLICATE,
            )
            await next(iter(tasks))
            replay = [decode(await outbound.get()).route_update for _ in range(3)]
            self.assertEqual(
                [update.update_id for update in replay],
                [update.update_id for update in first_updates],
            )
            self.assertEqual(manager.configure.call_count, 1)

            accepted_before_crash = new_envelope()
            accepted_before_crash.command.CopyFrom(command.command)
            accepted_before_crash.command.command_id = str(uuid4())
            acknowledgement = daemon._initial_acknowledgement(
                accepted_before_crash.command, set(), False
            )
            await asyncio.to_thread(
                daemon.store.accept_command,
                accepted_before_crash.command.command_id,
                encode_message(accepted_before_crash.command),
                encode_message(acknowledgement),
            )
            restarted = ArchiveClientDaemon(
                config,
                [probe, probe],
                [syncthing_probe],
                route_manager=manager,
            )
            resumed_outbound = asyncio.Queue()
            resumed_tasks = set()
            await restarted._resume_route_commands(
                resumed_outbound, resumed_tasks
            )
            await next(iter(resumed_tasks))
            resumed = [
                decode(await resumed_outbound.get()).route_update
                for _ in range(3)
            ]
            self.assertEqual([item.sequence for item in resumed], [1, 2, 3])
            self.assertEqual(manager.configure.call_count, 2)
            self.assertEqual(
                restarted._route_path("route-1"), root / "routes/route-1"
            )

    async def test_slow_inventory_does_not_block_heartbeat(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            service = ServiceConfig(
                "http://local", PurePosixPath("/api"), root
            )
            config = ClientConfig(
                "cache-1", "Cache 1", "cache", "ws://control", token,
                root / "state.db", root / "backups", service, service,
            )
            gate = threading.Event()

            def slow_inventory(_filter=""):
                gate.wait(2)
                return []

            reader = Mock(
                list_resources=Mock(side_effect=slow_inventory),
                get_resource=Mock(return_value=None),
            )
            probe = FilesystemProbe(root, True, True, True, True, True)
            daemon = ArchiveClientDaemon(
                config, [probe, probe], [], resource_reader=reader
            )
            await asyncio.to_thread(daemon.store.initialize)
            outbound = asyncio.Queue()
            tasks = set()
            command = new_envelope()
            command.command.command_id = str(uuid4())
            command.command.created_at.CopyFrom(command.sent_at)
            command.command.inventory_query.query_id = str(uuid4())
            command.command.inventory_query.scope = (
                inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES
            )
            await daemon._handle(command, outbound, tasks)
            self.assertEqual(
                decode(await outbound.get()).command_ack.status,
                control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
            )
            inventory_task = next(iter(tasks))
            heartbeat = new_envelope()
            heartbeat.heartbeat.sequence = 9
            await daemon._handle(heartbeat, outbound, tasks)
            self.assertEqual(
                decode(await outbound.get()).heartbeat_ack.sequence, 9
            )
            gate.set()
            await inventory_task
            self.assertTrue(
                decode(await outbound.get()).inventory_chunk.last_chunk
            )

    async def test_registration_heartbeat_and_duplicate_command(self):
        observed = {}
        job_id = str(uuid4())

        async def control(websocket):
            registration = decode(await websocket.recv())
            observed["token"] = registration.register_request.shared_token
            observed["root_names"] = [
                item.root_name
                for item in registration.register_request.capabilities.filesystems
            ]
            observed["addresses"] = list(
                registration.register_request.capabilities
                .syncthing_advertised_addresses
            )
            observed["device_id"] = (
                registration.register_request.capabilities.syncthing_device_id
            )
            observed["services"] = [
                item.service
                for item in registration.register_request.capabilities.services
            ]
            observed["features"] = list(
                registration.register_request.capabilities.features
            )
            response = new_envelope()
            response.correlation_id = registration.message_id
            response.register_response.status = client_pb2.REGISTRATION_STATUS_ACCEPTED
            response.register_response.negotiated_version.major = 1
            await websocket.send(encode(response))
            heartbeat = new_envelope()
            heartbeat.heartbeat.sequence = 7
            await websocket.send(encode(heartbeat))
            observed["heartbeat"] = decode(await websocket.recv()).heartbeat_ack.sequence
            command_id = str(uuid4())
            command = new_envelope()
            command.command.command_id = command_id
            command.command.created_at.CopyFrom(command.sent_at)
            command.command.request_job_snapshot.job_ids.append(job_id)
            await websocket.send(encode(command))
            observed["first"] = decode(await websocket.recv()).command_ack.status
            snapshot = decode(await websocket.recv())
            observed["snapshot"] = snapshot.WhichOneof("payload")
            observed["snapshot_job_id"] = (
                snapshot.job_snapshot.job.definition.job_id
            )
            duplicate = new_envelope()
            duplicate.command.CopyFrom(command.command)
            await websocket.send(encode(duplicate))
            observed["second"] = decode(await websocket.recv()).command_ack.status
            observed["duplicate_snapshot"] = (
                decode(await websocket.recv()).WhichOneof("payload")
            )
            unsupported = new_envelope()
            unsupported.command.command_id = str(uuid4())
            unsupported.command.created_at.CopyFrom(unsupported.sent_at)
            unsupported.command.assign_job.SetInParent()
            await websocket.send(encode(unsupported))
            rejected = decode(await websocket.recv()).command_ack
            observed["rejected"] = (rejected.status, rejected.error.code)
            unsupported_duplicate = new_envelope()
            unsupported_duplicate.command.CopyFrom(unsupported.command)
            await websocket.send(encode(unsupported_duplicate))
            observed["rejected_duplicate"] = (
                decode(await websocket.recv()).command_ack.status
            )
            inventory = new_envelope()
            inventory.command.command_id = str(uuid4())
            inventory.command.created_at.CopyFrom(inventory.sent_at)
            inventory.command.inventory_query.query_id = str(uuid4())
            inventory.command.inventory_query.scope = (
                inventory_pb2.INVENTORY_SCOPE_RESOURCE_SUMMARIES
            )
            await websocket.send(encode(inventory))
            observed["inventory_ack"] = (
                decode(await websocket.recv()).command_ack.status
            )
            result = decode(await websocket.recv()).inventory_chunk
            observed["inventory_result"] = (
                result.last_chunk, result.WhichOneof("payload")
            )

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            token = root / "token"
            token.write_text("shared-secret", encoding="utf-8")
            os.chmod(token, 0o600)
            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
                port = server.sockets[0].getsockname()[1]
                service = ServiceConfig(
                    "http://local", PurePosixPath("/api"), root,
                    advertised_addresses=("dynamic",),
                )
                config = ClientConfig(
                    "cache-1", "Cache 1", "cache",
                    f"ws://127.0.0.1:{port}", token,
                    root / "state.db", root / "backups", service, service,
                    ConnectionConfig(registration_timeout=2),
                )
                probe = FilesystemProbe(root, True, True, True, True, True)
                service_probe = ServiceProbe(
                    "syncthing", common_pb2.HEALTH_STATE_HEALTHY,
                    datetime.now(timezone.utc), version="v2", device_id="DEVICE",
                )
                daemon = ArchiveClientDaemon(
                    config, [probe, probe], [service_probe],
                    resource_reader=Mock(
                        list_resources=Mock(return_value=[]),
                        get_resource=Mock(return_value=None),
                    ),
                )
                await asyncio.to_thread(daemon.store.initialize)
                definition = job_pb2.JobDefinition(
                    job_id=job_id,
                    operation=job_pb2.JOB_OPERATION_ARCHIVE,
                )
                definition.created_at.GetCurrentTime()
                await asyncio.to_thread(
                    daemon.store.save_job,
                    job_id,
                    encode_message(definition),
                    "JOB_STATE_WAITING",
                    2,
                    3,
                    False,
                )
                await daemon._connection()

        self.assertEqual(observed["token"], "shared-secret")
        self.assertEqual(observed["root_names"], ["qbittorrent", "syncthing"])
        self.assertEqual(observed["addresses"], ["dynamic"])
        self.assertEqual(observed["device_id"], "DEVICE")
        self.assertEqual(observed["services"], ["syncthing"])
        self.assertIn(
            client_pb2.CLIENT_FEATURE_INVENTORY_CHUNKS,
            observed["features"],
        )
        self.assertIn(
            client_pb2.CLIENT_FEATURE_ROUTE_PROVISIONING,
            observed["features"],
        )
        self.assertEqual(observed["heartbeat"], 7)
        self.assertEqual(observed["first"], control_pb2.COMMAND_ACK_STATUS_ACCEPTED)
        self.assertEqual(observed["second"], control_pb2.COMMAND_ACK_STATUS_DUPLICATE)
        self.assertEqual(observed["snapshot"], "job_snapshot")
        self.assertEqual(observed["snapshot_job_id"], job_id)
        self.assertEqual(observed["duplicate_snapshot"], "job_snapshot")
        self.assertEqual(
            observed["rejected"],
            (
                control_pb2.COMMAND_ACK_STATUS_REJECTED,
                common_pb2.ERROR_CODE_INVALID_ARGUMENT,
            ),
        )
        self.assertEqual(
            observed["rejected_duplicate"],
            control_pb2.COMMAND_ACK_STATUS_REJECTED,
        )
        self.assertEqual(
            observed["inventory_ack"],
            control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
        )
        self.assertEqual(
            observed["inventory_result"], (True, "resource_summaries")
        )


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