"""In-process WebSocket tests for transfer_session push (incl. batch TTL expiry).

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

from __future__ import annotations

import asyncio
import json
from os.path import join
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import FastAPI, WebSocket
from starlette.testclient import TestClient

import mcp_proxy_adapter.api.handlers as handlers_mod
from mcp_proxy_adapter.api.handlers import expire_transfer_sessions_and_notify
from mcp_proxy_adapter.api.ws_endpoint import handle_websocket_job_push
from mcp_proxy_adapter.core.job_push.events import (
    ACTION_SUBSCRIBE,
    EVENT_TRANSFER_EXPIRED,
    SUBSCRIPTION_TYPE_TRANSFER_SESSION,
)
from mcp_proxy_adapter.core.job_push.subscriptions import (
    Subscriber,
    SubscriberRegistry,
    SubscriptionCategory,
)
from mcp_proxy_adapter.core.job_push.transfer_push import (
    set_transfer_push_app,
)
from mcp_proxy_adapter.transfer import TransferServerStore


def _merge_dict_for_tests(storage: str) -> dict[str, object]:
    return {
        "storage_dir": storage,
        "session_ttl_seconds": 3600,
        "max_file_size_bytes": 10_000_000,
        "default_chunk_size_bytes": 1024,
        "max_chunk_size_bytes": 1024,
        "checksum_algorithm": "sha256",
        "server_gzip_download_buffer_dir": None,
        "ws_push_enabled": True,
    }


def _make_store(tmp_storage: str) -> TransferServerStore:
    return TransferServerStore(
        storage_dir=tmp_storage,
        session_ttl_seconds=3600,
        max_file_size_bytes=10_000_000,
        default_chunk_size_bytes=1024,
        max_chunk_size_bytes=1024,
    )


def _patch_expires_to_past(store: TransferServerStore, transfer_id: str) -> None:
    path = join(store._sessions_dir(), f"{transfer_id}.json")
    with open(path, encoding="utf-8") as f:
        raw = json.load(f)
    raw["expires_at"] = "2000-01-01T00:00:00+00:00"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(raw, f, indent=2)


def test_transfer_expired_emitted_from_batch_expiry_path(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """TTL batch expiry pushes transfer_expired on the same portal as /ws (Starlette TestClient)."""
    storage = str(tmp_path / "tstore")
    store = _make_store(storage)
    receipt = store.create_upload_session(
        filename="f.txt",
        size_bytes=10,
        checksum_algorithm="sha256",
        checksum_value="0" * 64,
        compression="identity",
    )
    tid = receipt.transfer_id

    monkeypatch.setattr(handlers_mod, "_CACHED_TRANSFER_STORE", None)
    monkeypatch.setattr(handlers_mod, "get_transfer_store", lambda: store)
    monkeypatch.setattr(
        handlers_mod,
        "_merge_transfer_config",
        lambda: _merge_dict_for_tests(storage),
    )

    app = FastAPI()
    registry = SubscriberRegistry()
    app.state.job_push_registry = registry
    manager = MagicMock()
    manager.flush_replay_to_subscriber = AsyncMock()
    app.state.ws_delivery_manager = manager
    set_transfer_push_app(app)

    @app.websocket("/ws")
    async def _ws(websocket: WebSocket) -> None:
        await handle_websocket_job_push(websocket, app)

    try:
        with patch(
            "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
            return_value=True,
        ):
            with TestClient(app) as client:
                with client.websocket_connect("/ws") as ws:
                    ws.send_json(
                        {
                            "action": ACTION_SUBSCRIBE,
                            "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                            "transfer_id": tid,
                        }
                    )

                    # Let the /ws handler finish subscribe before we age out metadata (ordering).
                    async def _yield_to_ws() -> None:
                        await asyncio.sleep(0.05)

                    ws.portal.call(_yield_to_ws)
                    _patch_expires_to_past(store, tid)

                    async def _run_expire() -> None:
                        await expire_transfer_sessions_and_notify(store)

                    ws.portal.call(_run_expire)
                    msg = ws.receive_json()
    finally:
        set_transfer_push_app(None)

    assert msg.get("event") == EVENT_TRANSFER_EXPIRED, msg
    assert msg["subscription_type"] == SUBSCRIPTION_TYPE_TRANSFER_SESSION
    assert msg["transfer_id"] == tid
    assert msg["direction"] == "upload"
    assert msg["error_type"] == "TransferExpiredError"
    assert msg["progress_percent"] == 0
    assert "expired" in str(msg.get("message", "")).lower()


class _CaptureSubscriber(Subscriber):
    def __init__(self) -> None:
        self.messages: list[dict] = []

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


@pytest.mark.asyncio
async def test_batch_expiry_respects_ws_push_disabled(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """When ws_push_enabled is false, notify_transfer_state_changed no-ops (no capture)."""
    storage = str(tmp_path / "tstore2")
    store = _make_store(storage)
    receipt = store.create_upload_session(
        filename="g.bin",
        size_bytes=8,
        checksum_algorithm="sha256",
        checksum_value="1" * 64,
        compression="identity",
    )
    tid = receipt.transfer_id
    _patch_expires_to_past(store, tid)

    monkeypatch.setattr(handlers_mod, "_CACHED_TRANSFER_STORE", None)
    monkeypatch.setattr(handlers_mod, "get_transfer_store", lambda: store)
    monkeypatch.setattr(
        handlers_mod,
        "_merge_transfer_config",
        lambda: {**_merge_dict_for_tests(storage), "ws_push_enabled": False},
    )

    app = FastAPI()
    registry = SubscriberRegistry()
    app.state.job_push_registry = registry
    set_transfer_push_app(app)
    cap = _CaptureSubscriber()
    await registry.subscribe(tid, SubscriptionCategory.TRANSFER_SESSION, cap)
    try:
        await expire_transfer_sessions_and_notify(store)
    finally:
        set_transfer_push_app(None)

    assert cap.messages == []
    loaded = store._load_session_raw(tid)
    assert loaded is not None
    assert loaded.status == "expired"
