"""Real local WebSocket/SDK round trips; no device, credential, or paid endpoint."""

import asyncio
from contextlib import asynccontextmanager
from copy import deepcopy
import json

import httpx
from langchain_core.messages import HumanMessage, ToolMessage
import pytest
import pytest_asyncio
from websockets.asyncio.server import serve

from artemis.llm.diagnostics import ACTIVE_ATTEMPT
from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
from artemis.llm.router import ModelEndpoint, ModelFactory
from artemis.services.llm import RobustChatModelWrapper


def text_output(text="OK"):
    return [
        {
            "id": "msg_test",
            "type": "message",
            "role": "assistant",
            "status": "completed",
            "content": [{"type": "output_text", "text": text, "annotations": [], "logprobs": []}],
        }
    ]


class Wire:
    def __init__(self):
        self.requests = []
        self.connections = []
        self.outputs = []
        self.reject_previous = False
        self.drop = False
        self.stall = False
        self.status = "completed"
        self.sent = asyncio.Event()
        self.gate = None

    async def handle(self, socket):
        self.connections.append(socket)
        async for raw in socket:
            body = json.loads(raw)
            self.requests.append(body)
            self.sent.set()
            if self.reject_previous and body.get("previous_response_id"):
                await socket.send(
                    json.dumps(
                        {
                            "type": "error",
                            "status": 400,
                            "error": {"code": "previous_response_not_found"},
                        }
                    )
                )
                continue
            response = {
                "id": f"resp_{len(self.requests)}",
                "object": "response",
                "created_at": 1,
                "status": "in_progress",
                "model": "test-model",
                "output": [],
            }
            await socket.send(
                json.dumps({"type": "response.created", "sequence_number": 0, "response": response})
            )
            if self.drop:
                await socket.close()
                continue
            if self.stall:
                await socket.wait_closed()
                return
            if self.gate:
                await self.gate.wait()
            output = self.outputs.pop(0) if self.outputs else text_output()
            seq = 0
            for index, item in enumerate(output):
                seq += 1
                await socket.send(
                    json.dumps(
                        {
                            "type": "response.output_item.added",
                            "sequence_number": seq,
                            "output_index": index,
                            "item": {**item, "content": [], "arguments": ""},
                        }
                    )
                )
                for content_index, block in enumerate(item.get("content", [])):
                    if block.get("type") != "output_text":
                        continue
                    seq += 1
                    await socket.send(
                        json.dumps(
                            {
                                "type": "response.output_text.delta",
                                "sequence_number": seq,
                                "output_index": index,
                                "content_index": content_index,
                                "item_id": item["id"],
                                "delta": block["text"],
                            }
                        )
                    )
                    seq += 1
                    await socket.send(
                        json.dumps(
                            {
                                "type": "response.output_text.done",
                                "sequence_number": seq,
                                "output_index": index,
                                "content_index": content_index,
                                "item_id": item["id"],
                                "text": block["text"],
                            }
                        )
                    )
                seq += 1
                await socket.send(
                    json.dumps(
                        {
                            "type": "response.output_item.done",
                            "sequence_number": seq,
                            "output_index": index,
                            "item": item,
                        }
                    )
                )
            response.update(
                status=self.status,
                output=output,
                usage={
                    "input_tokens": 20,
                    "output_tokens": 3,
                    "total_tokens": 23,
                    "input_tokens_details": {"cached_tokens": 10},
                },
            )
            await socket.send(
                json.dumps(
                    {
                        "type": f"response.{self.status}",
                        "sequence_number": seq + 1,
                        "response": response,
                    }
                )
            )


@pytest_asyncio.fixture
async def wire(monkeypatch):
    monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False)
    instance = Wire()
    async with serve(instance.handle, "127.0.0.1", 0) as server:
        instance.base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1"
        yield instance


@asynccontextmanager
async def model(wire, **kwargs):
    endpoint = ModelEndpoint(
        model_name="test-model",
        api_key="test-secret",
        api_base=wire.base,
        transport="websocket",
        max_retries=0,
        timeout_seconds=2,
        **kwargs,
    )
    value = ModelFactory.create_model(endpoint)
    try:
        yield value
    finally:
        await value.http_async_client.aclose()
        value.http_client.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_sdk_roundtrip_incremental_history_and_usage(wire, stream):
    async with model(wire) as raw:
        messages = [
            HumanMessage(
                content=[
                    {"type": "text", "text": "Inspect"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/png;base64," + "A" * 80000},
                    },
                ]
            )
        ]

        async def call(history):
            if not stream:
                return await raw.ainvoke(history, prompt_cache_key="task:operator")
            chunks = [
                chunk async for chunk in raw.astream(history, prompt_cache_key="task:operator")
            ]
            return sum(chunks[1:], chunks[0])

        first = await call(messages)
        ensure_response_complete(first)
        second = await call([*messages, first, HumanMessage(content="Continue")])
        ensure_response_complete(second)
        assert second.usage_metadata["input_token_details"]["cache_read"] == 10
        assert wire.requests[1].get("previous_response_id") == "resp_1", {
            "first": first.content,
            "tail": wire.requests[1]["input"][1:],
        }
        assert wire.requests[1]["input"] == [
            {"type": "message", "role": "user", "content": "Continue"}
        ]
        assert "stream" not in wire.requests[0]
        assert wire.requests[0]["type"] == "response.create"
        assert wire.requests[0]["store"] is False
        assert len(json.dumps(wire.requests[1])) < len(json.dumps(wire.requests[0])) / 100
        assert len(wire.connections) == 1
        assert wire.connections[0].request.path == "/v1/responses"
        assert wire.connections[0].request.headers["Authorization"] == "Bearer test-secret"


@pytest.mark.asyncio
async def test_tool_call_chain_preserves_call_id(wire):
    wire.outputs.append(
        [
            {
                "id": "fc_test",
                "type": "function_call",
                "call_id": "call_test",
                "name": "lookup",
                "arguments": '{"x":1}',
                "status": "completed",
            }
        ]
    )
    async with model(wire) as raw:
        messages = [HumanMessage(content="Look up")]
        first = await raw.ainvoke(messages, prompt_cache_key="task:operator")
        await raw.ainvoke(
            [*messages, first, ToolMessage(content="found", tool_call_id="call_test")],
            prompt_cache_key="task:operator",
        )
        assert wire.requests[1]["previous_response_id"] == "resp_1"
        assert wire.requests[1]["input"] == [
            {"type": "function_call_output", "call_id": "call_test", "output": "found"}
        ]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "change", ["history", "options", "scope", "new_connection", "server_closed"]
)
async def test_chain_resets_on_context_or_connection_change(wire, change):
    async with model(wire) as raw:
        messages = [HumanMessage(content="Original")]
        first = await raw.ainvoke(messages, prompt_cache_key="task:operator")
        history = [*messages, first, HumanMessage(content="Next")]
        options = {"prompt_cache_key": "task:operator"}
        if change == "history":
            history[0] = HumanMessage(content="Compacted")
        if change == "options":
            options["temperature"] = 0.3
        if change == "scope":
            options["prompt_cache_key"] = "task:checker"
        if change == "new_connection":
            await raw.http_async_client._transport.sessions[
                next(iter(raw.http_async_client._transport.sessions))
            ].close()
        if change == "server_closed":
            await wire.connections[0].close()
            await wire.connections[0].wait_closed()
        await raw.ainvoke(history, **options)
        assert "previous_response_id" not in wire.requests[1]
        assert len(wire.requests[1]["input"]) == 3


@pytest.mark.asyncio
async def test_uncached_id_recovers_once_with_full_context(wire):
    async with model(wire) as raw:
        history = [HumanMessage(content="First")]
        first = await raw.ainvoke(history, prompt_cache_key="task:operator")
        wire.reject_previous = True
        result = await raw.ainvoke(
            [*history, first, HumanMessage(content="Next")], prompt_cache_key="task:operator"
        )
        ensure_response_complete(result)
        assert len(wire.requests) == 3
        assert wire.requests[1]["previous_response_id"] == "resp_1"
        assert "previous_response_id" not in wire.requests[2]
        assert len(wire.requests[2]["input"]) == 3


@pytest.mark.asyncio
async def test_missing_affinity_never_shares_connections(wire):
    async with model(wire) as raw:
        history = [HumanMessage(content="First")]
        first = await raw.ainvoke(history)
        await raw.ainvoke([*history, first, HumanMessage(content="Next")])
        assert len(wire.connections) == 2
        assert "previous_response_id" not in wire.requests[1]


@pytest.mark.asyncio
async def test_cancel_closes_socket_and_next_call_starts_full(wire):
    async with model(wire) as raw:
        wire.stall = True
        task = asyncio.create_task(
            raw.ainvoke([HumanMessage(content="Wait")], prompt_cache_key="task:operator")
        )
        await asyncio.wait_for(wire.sent.wait(), 2)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        await asyncio.wait_for(wire.connections[0].wait_closed(), 2)
        wire.stall = False
        await raw.ainvoke([HumanMessage(content="Next")], prompt_cache_key="task:operator")
        assert len(wire.connections) == 2
        assert "previous_response_id" not in wire.requests[1]


@pytest.mark.asyncio
async def test_disconnect_never_silently_resends_model_request(wire):
    async with model(wire) as raw:
        wire.drop = True
        with pytest.raises(Exception, match="connection|Connection"):
            await raw.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator")
        assert len(wire.requests) == 1


@pytest.mark.asyncio
async def test_incomplete_result_does_not_advance_chain(wire):
    async with model(wire) as raw:
        wire.status = "incomplete"
        first = await raw.ainvoke([HumanMessage(content="First")], prompt_cache_key="task:operator")
        with pytest.raises(ResponseIntegrityError):
            ensure_response_complete(first)
        wire.status = "completed"
        await raw.ainvoke(
            [HumanMessage(content="First"), first, HumanMessage(content="Next")],
            prompt_cache_key="task:operator",
        )
        assert "previous_response_id" not in wire.requests[1]


@pytest.mark.asyncio
async def test_diagnostics_measure_actual_wire_payload_without_logging_content(wire):
    class Attempt:
        data = {}

        def publish_if_due(self):
            pass

    attempt = Attempt()
    async with model(wire) as raw:
        # Disable the unrelated HTTP trace hook: this checks the transport's
        # own byte counters, alongside the real SDK request serialization.
        raw.http_async_client.event_hooks = {"request": [], "response": []}
        first = await raw.ainvoke(
            [HumanMessage(content="A" * 5000)], prompt_cache_key="task:operator"
        )
        token = ACTIVE_ATTEMPT.set(attempt)
        try:
            await raw.ainvoke(
                [HumanMessage(content="A" * 5000), first, HumanMessage(content="Next")],
                prompt_cache_key="task:operator",
            )
        finally:
            ACTIVE_ATTEMPT.reset(token)
        stats = deepcopy(attempt.data["transport"])
        assert stats["incremental"] is True
        assert stats["sent_payload_bytes"] == len(
            json.dumps(
                wire.requests[1], ensure_ascii=False, separators=(",", ":"), sort_keys=True
            ).encode()
        )
        assert stats["payload_bytes_saved"] > 4000
        assert "AAAA" not in json.dumps(stats)


@pytest.mark.asyncio
async def test_gateway_rejects_upgrade_with_explicit_error(wire):
    import openai
    from websockets.http11 import Response
    from websockets.datastructures import Headers

    async with serve(
        wire.handle,
        "127.0.0.1",
        0,
        process_request=lambda *_: Response(404, "Not Found", Headers(), b""),
    ) as server:
        wire.base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1"
        async with model(wire) as raw:
            with pytest.raises(openai.NotFoundError, match="WebSocket"):
                await raw.ainvoke([HumanMessage(content="Test")])


@pytest.mark.asyncio
async def test_wrapper_retains_stream_guards_and_diagnostic_events(wire, monkeypatch):
    from artemis.services import llm
    from tests.unit.test_responses_wire_lifecycle import Engine

    engine = Engine()
    monkeypatch.setattr(llm, "_get_current_data_engine", lambda: engine)
    async with model(wire) as raw:
        wrapper = RobustChatModelWrapper(
            raw,
            endpoint=ModelEndpoint(model_name="test-model", transport="websocket", max_retries=0),
        )
        result = await wrapper.ainvoke(
            [HumanMessage(content="Test")], prompt_cache_key="task:operator"
        )
        ensure_response_complete(result)
        assert engine.data["transport"]["mode"] == "websocket"
        assert engine.data["transport"]["received_payload_bytes"] > 0
        assert engine.data.get("http_status") is None


@pytest.mark.asyncio
async def test_parallel_scopes_keep_independent_response_chains(wire):
    async with model(wire) as raw:

        async def conversation(scope):
            history = [HumanMessage(content=f"First {scope}")]
            first = await raw.ainvoke(history, prompt_cache_key=scope)
            await raw.ainvoke(
                [*history, first, HumanMessage(content=f"Next {scope}")], prompt_cache_key=scope
            )
            return scope, first.response_metadata["id"]

        results = await asyncio.gather(
            conversation("task-a:operator"), conversation("task-b:operator")
        )
        assert len(wire.connections) == 2
        for scope, response_id in results:
            calls = [item for item in wire.requests if item["prompt_cache_key"] == scope]
            assert len(calls) == 2
            assert calls[1]["previous_response_id"] == response_id
            assert len(calls[1]["input"]) == 1


@pytest.mark.asyncio
async def test_idle_connections_are_closed_and_removed(wire, monkeypatch):
    from artemis.llm import websocket

    monkeypatch.setattr(websocket, "IDLE_SECONDS", 0.01)
    async with model(wire) as raw:
        await raw.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator")
        await asyncio.wait_for(wire.connections[0].wait_closed(), 2)
        assert not raw.http_async_client._transport.sessions


@pytest.mark.asyncio
async def test_cancel_waiting_turn_keeps_active_chain_and_idle_cleanup(wire, monkeypatch):
    from artemis.llm import websocket

    # Only shorten the final idle interval; scheduler latency must not break the first chain.
    monkeypatch.setattr(websocket, "IDLE_SECONDS", 60)
    wire.gate = asyncio.Event()
    async with model(wire) as raw:
        history = [HumanMessage(content="First")]
        first = asyncio.create_task(raw.ainvoke(history, prompt_cache_key="task:operator"))
        await asyncio.wait_for(wire.sent.wait(), 2)
        queued = asyncio.create_task(raw.ainvoke(history, prompt_cache_key="task:operator"))
        async with asyncio.timeout(2):
            while next(iter(raw.http_async_client._transport.sessions.values())).users != 2:
                await asyncio.sleep(0)
        queued.cancel()
        with pytest.raises(asyncio.CancelledError):
            await queued
        wire.gate.set()
        result = await first
        monkeypatch.setattr(websocket, "IDLE_SECONDS", 0.01)
        await raw.ainvoke(
            [*history, result, HumanMessage(content="Next")], prompt_cache_key="task:operator"
        )
        assert len(wire.connections) == 1
        assert len(wire.requests) == 2
        assert wire.requests[1]["previous_response_id"] == result.response_metadata["id"]
        await asyncio.wait_for(wire.connections[0].wait_closed(), 2)
        assert not raw.http_async_client._transport.sessions


@pytest.mark.asyncio
async def test_pool_eviction_does_not_split_concurrent_scope(wire, monkeypatch):
    from artemis.llm import websocket

    monkeypatch.setattr(websocket, "MAX_SESSIONS", 1)
    closing, resume = asyncio.Event(), asyncio.Event()
    async with model(wire) as raw:
        await raw.ainvoke([HumanMessage(content="Old")], prompt_cache_key="old")
        old = next(iter(raw.http_async_client._transport.sessions.values()))
        close = old.close

        async def delayed_close():
            closing.set()
            await resume.wait()
            await close()

        monkeypatch.setattr(old, "close", delayed_close)
        first = asyncio.create_task(
            raw.ainvoke([HumanMessage(content="New first")], prompt_cache_key="new")
        )
        try:
            await asyncio.wait_for(closing.wait(), 2)
            await asyncio.wait_for(
                asyncio.gather(
                    first,
                    raw.ainvoke([HumanMessage(content="New second")], prompt_cache_key="new"),
                ),
                2,
            )
            assert len(wire.connections) == 2
            assert len(wire.requests) == 3
        finally:
            resume.set()


@pytest.mark.asyncio
async def test_structured_output_uses_same_websocket_path(wire):
    from pydantic import BaseModel

    class Result(BaseModel):
        ok: bool

    wire.outputs.append(
        [
            {
                "id": "fc_test",
                "type": "function_call",
                "call_id": "call_test",
                "name": "Result",
                "arguments": '{"ok":true}',
                "status": "completed",
            }
        ]
    )
    async with model(wire) as raw:
        result = (
            await RobustChatModelWrapper(
                raw, endpoint=ModelEndpoint(transport="websocket", max_retries=0)
            )
            .with_structured_output(Result)
            .ainvoke([HumanMessage(content="Test")])
        )
        assert result.ok is True
        assert len(wire.requests) == 1


@pytest.mark.asyncio
async def test_invalid_tool_arguments_are_rejected_before_use(wire, monkeypatch):
    from artemis.services import llm
    from artemis.llm.reliability import LLMExhaustedError

    monkeypatch.setattr(llm, "_get_current_data_engine", lambda: None)
    wire.outputs.append(
        [
            {
                "id": "fc_test",
                "type": "function_call",
                "call_id": "call_test",
                "name": "lookup",
                "arguments": '{"x":',
                "status": "completed",
            }
        ]
    )
    async with model(wire) as raw:
        wrapper = RobustChatModelWrapper(
            raw, endpoint=ModelEndpoint(transport="websocket", max_retries=0)
        )
        with pytest.raises(LLMExhaustedError):
            await wrapper.ainvoke([HumanMessage(content="Test")], prompt_cache_key="task:operator")
        assert len(wire.requests) == 1


@pytest.mark.asyncio
async def test_redirect_does_not_forward_authorization(wire):
    from websockets.http11 import Response
    from websockets.datastructures import Headers

    async with serve(
        wire.handle,
        "127.0.0.1",
        0,
        process_request=lambda *_: Response(
            307,
            "Redirect",
            Headers({"Location": wire.base.replace("http://", "ws://") + "/responses"}),
            b"",
        ),
    ) as server:
        redirect_base = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}/v1"
        original_base, wire.base = wire.base, redirect_base
        try:
            async with model(wire) as raw:
                with pytest.raises(Exception, match="redirects are disabled"):
                    await raw.ainvoke([HumanMessage(content="Test")])
            assert not wire.connections
        finally:
            wire.base = original_base