"""Concurrency and race tests: subscribe-after-acceptance, replay-after-completion.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import asyncio
from typing import Dict, List

import pytest
import httpx
from fastapi import FastAPI

from mcp_proxy_adapter.core.job_push.subscriptions import (
    SubscriptionCategory,
)


class CollectingSubscriber:
    """Subscriber that appends received payloads to a list."""

    def __init__(self) -> None:
        self.payloads: List[Dict[str, Any]] = []

    async def send(self, payload: dict) -> None:
        self.payloads.append(dict(payload))


@pytest.mark.asyncio
async def test_subscribe_after_acceptance_receives_exactly_one_terminal(
    execute_async_test_app: FastAPI,
    async_asgi_client: httpx.AsyncClient,
    jsonrpc_execute_async_url: str,
) -> None:
    """Subscribe with ws_id after HTTP acceptance returns; receive exactly one terminal."""
    response = await async_asgi_client.post(
        jsonrpc_execute_async_url,
        json={
            "jsonrpc": "2.0",
            "method": "execute_async",
            "params": {"command": "echo", "params": {"message": "after-accept"}},
            "id": 1,
        },
    )
    assert response.status_code == 200
    ws_id = response.json()["result"]["ws_id"]
    assert ws_id

    registry = execute_async_test_app.state.ws_delivery_manager._registry
    sub = CollectingSubscriber()
    await registry.subscribe(ws_id, SubscriptionCategory.WS_DELIVERY, sub)

    await asyncio.sleep(0.4)
    assert len(sub.payloads) >= 1
    terminal_events = [p for p in sub.payloads if p.get("event") in ("job_completed", "delivery_completed")]
    assert len(terminal_events) == 1


@pytest.mark.asyncio
async def test_completion_before_subscribe_replay_delivers_terminal(
    execute_async_test_app: FastAPI,
    async_asgi_client: httpx.AsyncClient,
    jsonrpc_execute_async_url: str,
) -> None:
    """Complete command before subscribe; then subscribe; replay delivers terminal."""
    response = await async_asgi_client.post(
        jsonrpc_execute_async_url,
        json={
            "jsonrpc": "2.0",
            "method": "execute_async",
            "params": {"command": "echo", "params": {"message": "replay-test"}},
            "id": 1,
        },
    )
    assert response.status_code == 200
    ws_id = response.json()["result"]["ws_id"]
    assert ws_id

    await asyncio.sleep(0.4)
    manager = execute_async_test_app.state.ws_delivery_manager
    sub = CollectingSubscriber()
    await manager.flush_replay_to_subscriber(ws_id, sub)

    assert len(sub.payloads) == 1
    assert sub.payloads[0].get("event") in ("job_completed", "delivery_completed")


@pytest.mark.asyncio
async def test_subscribe_before_request_does_not_violate_contract(
    execute_async_test_app: FastAPI,
    async_asgi_client: httpx.AsyncClient,
    jsonrpc_execute_async_url: str,
) -> None:
    """Early subscribe with wrong/empty id must not receive execute_async completion."""
    registry = execute_async_test_app.state.ws_delivery_manager._registry
    sub = CollectingSubscriber()
    await registry.subscribe("wrong-ws-id-never-issued", SubscriptionCategory.WS_DELIVERY, sub)

    response = await async_asgi_client.post(
        jsonrpc_execute_async_url,
        json={
            "jsonrpc": "2.0",
            "method": "execute_async",
            "params": {"command": "echo", "params": {"message": "other"}},
            "id": 1,
        },
    )
    assert response.status_code == 200
    ws_id = response.json()["result"]["ws_id"]
    assert ws_id != "wrong-ws-id-never-issued"

    await asyncio.sleep(0.4)
    assert len(sub.payloads) == 0


@pytest.mark.asyncio
async def test_interleaved_poll_and_async_terminal_no_duplicate_terminal(
    async_asgi_client: httpx.AsyncClient,
    jsonrpc_execute_async_url: str,
) -> None:
    """Two execute_async calls each yield one terminal (no cross-delivery)."""
    responses = await asyncio.gather(
        async_asgi_client.post(
            jsonrpc_execute_async_url,
            json={
                "jsonrpc": "2.0",
                "method": "execute_async",
                "params": {"command": "echo", "params": {"message": "a"}},
                "id": 1,
            },
        ),
        async_asgi_client.post(
            jsonrpc_execute_async_url,
            json={
                "jsonrpc": "2.0",
                "method": "execute_async",
                "params": {"command": "echo", "params": {"message": "b"}},
                "id": 2,
            },
        ),
    )
    ws_ids = [r.json()["result"]["ws_id"] for r in responses]
    assert len(ws_ids) == 2
    assert ws_ids[0] != ws_ids[1]
    await asyncio.sleep(0.5)
