from __future__ import annotations as _annotations

import asyncio

import pytest
from acp import PROTOCOL_VERSION
from acp.exceptions import RequestError
from pydantic_acp import __version__ as pydantic_acp_version

from .support import (
    UTC,
    AcpSessionContext,
    AdapterConfig,
    AdapterModel,
    AdapterPromptCapabilities,
    Agent,
    AgentBridgeBuilder,
    AgentFactory,
    AgentMessageChunk,
    AgentSource,
    AvailableCommandsUpdate,
    ClientFilesystemBackend,
    ClientHostContext,
    ClientTerminalBackend,
    FactoryAgentSource,
    FileSessionStore,
    FilesystemBackend,
    McpBridge,
    McpServerDefinition,
    MemorySessionStore,
    ModelSelectionState,
    ModeState,
    Path,
    PrepareToolsBridge,
    PrepareToolsMode,
    RecordingClient,
    SessionInfoUpdate,
    StaticAgentSource,
    TerminalBackend,
    TestModel,
    ToolCallProgress,
    ToolCallStart,
    UserMessageChunk,
    agent_message_texts,
    create_acp_agent,
    datetime,
    text_block,
)


def test_initialize_uses_configured_prompt_capabilities(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="unused")),
        config=AdapterConfig(
            prompt_capabilities=AdapterPromptCapabilities(
                audio=False,
                image=False,
                embedded_context=False,
            ),
            session_store=MemorySessionStore(),
        ),
    )

    response = asyncio.run(adapter.initialize(protocol_version=PROTOCOL_VERSION))

    assert response.agent_capabilities is not None
    prompt_capabilities = response.agent_capabilities.prompt_capabilities
    assert prompt_capabilities is not None
    assert prompt_capabilities.audio is False
    assert prompt_capabilities.image is False
    assert prompt_capabilities.embedded_context is False


def test_prompt_and_load_session_replay_history(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="Hello from ACP")),
        config=AdapterConfig(
            agent_name="pydantic-acp",
            agent_title="Pydantic ACP",
            agent_version="0.1.0",
            session_store=MemorySessionStore(),
        ),
    )
    client = RecordingClient()
    adapter.on_connect(client)

    initialize_response = asyncio.run(adapter.initialize(protocol_version=1))
    assert initialize_response.agent_info is not None
    assert initialize_response.agent_info.name == "pydantic-acp"
    assert initialize_response.agent_capabilities is not None
    assert initialize_response.agent_capabilities.load_session is True
    assert initialize_response.agent_capabilities.prompt_capabilities is not None
    assert initialize_response.agent_capabilities.prompt_capabilities.audio is True
    assert initialize_response.agent_capabilities.prompt_capabilities.image is True
    assert initialize_response.agent_capabilities.prompt_capabilities.embedded_context is True

    new_session_response = asyncio.run(adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    prompt_response = asyncio.run(
        adapter.prompt(
            prompt=[text_block("Summarize the change.")],
            session_id=new_session_response.session_id,
            message_id="user-message-1",
        ),
    )

    assert prompt_response.stop_reason == "end_turn"
    assert prompt_response.user_message_id == "user-message-1"
    assert all(session_id == new_session_response.session_id for session_id, _ in client.updates)
    agent_updates = [
        update for _, update in client.updates if isinstance(update, AgentMessageChunk)
    ]
    assert len(agent_updates) >= 2
    assert {update.message_id for update in agent_updates} == {agent_updates[0].message_id}
    assert agent_message_texts(client) == ["Hello from ACP"]

    client.updates.clear()
    load_response = asyncio.run(
        adapter.load_session(
            cwd=str(tmp_path),
            session_id=new_session_response.session_id,
            mcp_servers=[],
        ),
    )

    assert load_response is not None
    replayed_update_types = [
        type(update)
        for _, update in client.updates
        if not isinstance(update, AvailableCommandsUpdate | SessionInfoUpdate)
    ]
    assert replayed_update_types[0] is UserMessageChunk
    assert replayed_update_types[1:] == [AgentMessageChunk] * (len(replayed_update_types) - 1)
    session_info = next(
        update for _, update in client.updates if isinstance(update, SessionInfoUpdate)
    )
    assert session_info.title == "Summarize the change."

    user_update = client.updates[0][1]
    assert isinstance(user_update, UserMessageChunk)
    assert user_update.content.text == "Summarize the change."
    assert agent_message_texts(client) == ["Hello from ACP"]


def test_initialize_uses_static_agent_name_by_default() -> None:
    adapter = create_acp_agent(agent=Agent(TestModel(), name="demo-agent-name"))

    initialize_response = asyncio.run(adapter.initialize(protocol_version=1))

    assert initialize_response.agent_info is not None
    assert initialize_response.agent_info.name == "demo-agent-name"
    assert initialize_response.agent_info.title == "Pydantic ACP"
    assert initialize_response.agent_info.version == pydantic_acp_version


def test_initialize_negotiates_protocol_and_exposes_mcp_capabilities() -> None:
    mcp_bridge = McpBridge(
        servers=[
            McpServerDefinition(server_id="http-server", name="HTTP", transport="http"),
            McpServerDefinition(server_id="sse-server", name="SSE", transport="sse"),
        ],
    )
    adapter = create_acp_agent(
        agent=Agent(TestModel()),
        config=AdapterConfig(
            agent_name="configured-agent",
            agent_title="Configured ACP",
            agent_version="9.9.9",
            capability_bridges=[mcp_bridge],
        ),
    )

    initialize_response = asyncio.run(adapter.initialize(protocol_version=PROTOCOL_VERSION + 10))

    assert initialize_response.protocol_version == PROTOCOL_VERSION
    assert initialize_response.agent_info is not None
    assert initialize_response.agent_info.name == "configured-agent"
    assert initialize_response.agent_info.title == "Configured ACP"
    assert initialize_response.agent_info.version == "9.9.9"
    assert initialize_response.agent_capabilities is not None
    assert initialize_response.agent_capabilities.mcp_capabilities is not None
    assert initialize_response.agent_capabilities.mcp_capabilities.http is True
    assert initialize_response.agent_capabilities.mcp_capabilities.sse is True
    assert initialize_response.agent_capabilities.session_capabilities is not None
    assert initialize_response.agent_capabilities.session_capabilities.fork is not None
    assert initialize_response.agent_capabilities.session_capabilities.resume is not None


def test_authenticate_cancel_and_extension_methods_follow_public_contract() -> None:
    adapter = create_acp_agent(agent=Agent(TestModel()))

    assert asyncio.run(adapter.authenticate("demo-method")) is None
    assert asyncio.run(adapter.cancel(session_id="missing-session")) is None
    assert asyncio.run(adapter.ext_notification("demo.notify", {"value": 1})) is None
    with pytest.raises(RequestError):
        asyncio.run(adapter.ext_method("demo.unknown", {"value": 1}))


def test_prompt_projects_tool_calls(tmp_path: Path) -> None:
    tool_model = TestModel(call_tools=["read_file"], custom_output_text="done")
    agent = Agent(tool_model)

    @agent.tool_plain
    def read_file(path: str) -> str:
        return f"contents:{path}"

    adapter = create_acp_agent(
        agent=agent,
        config=AdapterConfig(session_store=MemorySessionStore()),
    )
    client = RecordingClient()
    adapter.on_connect(client)

    new_session_response = asyncio.run(adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    asyncio.run(
        adapter.prompt(
            prompt=[text_block("Read the file.")],
            session_id=new_session_response.session_id,
        ),
    )

    tool_updates = [
        update
        for _, update in client.updates
        if isinstance(update, (ToolCallStart, ToolCallProgress))
    ]
    assert len(tool_updates) == 2

    tool_start = tool_updates[0]
    tool_completion = tool_updates[1]
    assert isinstance(tool_start, ToolCallStart)
    assert tool_start.title == "read_file"
    assert tool_start.kind == "read"
    assert tool_start.raw_input == {"path": "a"}
    assert isinstance(tool_completion, ToolCallProgress)
    assert tool_completion.status == "completed"
    assert tool_completion.raw_output == "contents:a"


def test_prompt_streams_text_chunks_with_shared_message_id(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="Hello from ACP")),
        config=AdapterConfig(session_store=MemorySessionStore()),
    )
    client = RecordingClient()
    adapter.on_connect(client)

    session = asyncio.run(adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    asyncio.run(
        adapter.prompt(
            prompt=[text_block("Stream the answer.")],
            session_id=session.session_id,
        ),
    )

    agent_updates = [
        update for _, update in client.updates if isinstance(update, AgentMessageChunk)
    ]
    assert len(agent_updates) >= 2
    assert all(update.message_id == agent_updates[0].message_id for update in agent_updates)
    assert agent_message_texts(client) == ["Hello from ACP"]


def test_optional_text_output_none_does_not_emit_literal_null(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text=""), output_type=str | None),
        config=AdapterConfig(session_store=MemorySessionStore()),
    )
    client = RecordingClient()
    adapter.on_connect(client)

    session = asyncio.run(adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    response = asyncio.run(
        adapter.prompt(
            prompt=[text_block("Return nothing if no answer is available.")],
            session_id=session.session_id,
        ),
    )

    assert response.stop_reason == "end_turn"
    assert agent_message_texts(client) == []


def test_file_session_store_round_trip(tmp_path: Path) -> None:
    store = FileSessionStore(tmp_path / "sessions")
    session = AcpSessionContext(
        session_id="session-123",
        cwd=tmp_path,
        created_at=datetime.now(UTC),
        updated_at=datetime.now(UTC),
        title="My Session",
        message_history_json='[{"role":"user"}]',
    )

    store.save(session)
    loaded_session = store.get("session-123")

    assert loaded_session is not None
    assert loaded_session.session_id == "session-123"
    assert loaded_session.title == "My Session"


def test_load_session_can_resume_with_persisted_mode_after_adapter_restart(
    tmp_path: Path,
) -> None:
    def ask_tools(ctx, tool_defs):
        del ctx  # pragma: no cover
        return list(tool_defs)  # pragma: no cover

    def review_tools(ctx, tool_defs):
        del ctx  # pragma: no cover
        return list(tool_defs)  # pragma: no cover

    store = FileSessionStore(tmp_path / "sessions")
    config = AdapterConfig(
        session_store=store,
        capability_bridges=[
            PrepareToolsBridge(
                default_mode_id="ask",
                modes=[
                    PrepareToolsMode(
                        id="ask",
                        name="Ask",
                        description="Default mode.",
                        prepare_func=ask_tools,
                    ),
                    PrepareToolsMode(
                        id="review",
                        name="Review",
                        description="Review mode.",
                        prepare_func=review_tools,
                    ),
                ],
            ),
        ],
    )
    first_adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="first-response")),
        config=config,
    )
    first_client = RecordingClient()
    first_adapter.on_connect(first_client)

    session = asyncio.run(first_adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    set_mode_response = asyncio.run(
        first_adapter.set_session_mode(mode_id="review", session_id=session.session_id),
    )
    assert set_mode_response is not None
    prompt_response = asyncio.run(
        first_adapter.prompt(
            prompt=[text_block("First run before restart.")],
            session_id=session.session_id,
        ),
    )
    assert prompt_response.stop_reason == "end_turn"

    restarted_adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="second-response")),
        config=AdapterConfig(
            session_store=store,
            capability_bridges=list(config.capability_bridges),
        ),
    )
    restarted_client = RecordingClient()
    restarted_adapter.on_connect(restarted_client)

    loaded = asyncio.run(
        restarted_adapter.load_session(
            cwd=str(tmp_path),
            session_id=session.session_id,
            mcp_servers=[],
        ),
    )
    assert loaded is not None
    assert loaded.modes is not None
    assert loaded.modes.current_mode_id == "review"
    assert agent_message_texts(restarted_client) == ["first-response"]

    follow_up = asyncio.run(
        restarted_adapter.prompt(
            prompt=[text_block("Continue after restart.")],
            session_id=session.session_id,
        ),
    )
    assert follow_up.stop_reason == "end_turn"
    assert agent_message_texts(restarted_client) == [
        "first-response",
        "second-response",
    ]


def test_file_session_store_skips_malformed_sessions_in_public_flows(
    tmp_path: Path,
) -> None:
    store = FileSessionStore(tmp_path / "sessions")
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="ok")),
        config=AdapterConfig(session_store=store),
    )

    valid_session = asyncio.run(adapter.new_session(cwd=str(tmp_path), mcp_servers=[]))
    (store.root / "broken.json").write_text("{not-json", encoding="utf-8")

    listed = asyncio.run(adapter.list_sessions())
    assert [session.session_id for session in listed.sessions] == [valid_session.session_id]

    loaded = asyncio.run(
        adapter.load_session(cwd=str(tmp_path), session_id="broken", mcp_servers=[]),
    )
    assert loaded is None


def test_interleaved_sessions_keep_selected_models_isolated(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="default")),
        config=AdapterConfig(
            session_store=MemorySessionStore(),
            available_models=[
                AdapterModel(
                    model_id="model-a",
                    name="Model A",
                    override=TestModel(custom_output_text="alpha"),
                ),
                AdapterModel(
                    model_id="model-b",
                    name="Model B",
                    override=TestModel(custom_output_text="beta"),
                ),
            ],
        ),
    )
    client = RecordingClient()
    adapter.on_connect(client)

    first = asyncio.run(adapter.new_session(cwd=str(tmp_path / "first"), mcp_servers=[]))
    second = asyncio.run(adapter.new_session(cwd=str(tmp_path / "second"), mcp_servers=[]))

    assert asyncio.run(adapter.set_session_model("model-a", first.session_id)) is not None
    assert asyncio.run(adapter.set_session_model("model-b", second.session_id)) is not None

    asyncio.run(adapter.prompt(prompt=[text_block("Run first once.")], session_id=first.session_id))
    asyncio.run(
        adapter.prompt(prompt=[text_block("Run second once.")], session_id=second.session_id),
    )
    asyncio.run(
        adapter.prompt(prompt=[text_block("Run first again.")], session_id=first.session_id),
    )

    first_updates = [
        update for session_id, update in client.updates if session_id == first.session_id
    ]
    second_updates = [
        update for session_id, update in client.updates if session_id == second.session_id
    ]
    first_messages: dict[str, str] = {}
    for update in first_updates:
        if isinstance(update, AgentMessageChunk):
            if update.message_id is None:  # pragma: no branch
                continue  # pragma: no cover
            first_messages[update.message_id] = (
                first_messages.get(update.message_id, "") + update.content.text
            )
    second_messages: dict[str, str] = {}
    for update in second_updates:
        if isinstance(update, AgentMessageChunk):
            if update.message_id is None:  # pragma: no branch
                continue  # pragma: no cover
            second_messages[update.message_id] = (
                second_messages.get(update.message_id, "") + update.content.text
            )

    assert list(first_messages.values()) == ["alpha", "alpha"]
    assert list(second_messages.values()) == ["beta"]


def test_list_sessions_filters_by_cwd_and_orders_latest_first(tmp_path: Path) -> None:
    adapter = create_acp_agent(
        agent=Agent(TestModel(custom_output_text="listed")),
        config=AdapterConfig(session_store=MemorySessionStore()),
    )

    alpha_session = asyncio.run(adapter.new_session(cwd=str(tmp_path / "alpha"), mcp_servers=[]))
    beta_session = asyncio.run(adapter.new_session(cwd=str(tmp_path / "beta"), mcp_servers=[]))
    asyncio.run(
        adapter.prompt(
            prompt=[text_block("Refresh alpha ordering.")],
            session_id=alpha_session.session_id,
        ),
    )

    listed_all = asyncio.run(adapter.list_sessions())
    listed_alpha = asyncio.run(adapter.list_sessions(cwd=str(tmp_path / "alpha")))

    assert [session.session_id for session in listed_all.sessions] == [
        alpha_session.session_id,
        beta_session.session_id,
    ]
    assert [session.cwd for session in listed_alpha.sessions] == [str(tmp_path / "alpha")]


def test_public_agent_source_exports_are_available() -> None:
    assert AgentFactory is not None
    assert AgentBridgeBuilder is not None
    assert AgentSource is not None
    assert ClientFilesystemBackend is not None
    assert ClientHostContext is not None
    assert ClientTerminalBackend is not None
    assert FilesystemBackend is not None
    assert FactoryAgentSource is not None
    assert StaticAgentSource is not None
    assert TerminalBackend is not None


def test_public_provider_exports_are_available() -> None:
    assert ModelSelectionState is not None
    assert ModeState is not None


def test_create_acp_agent_requires_exactly_one_source() -> None:
    with pytest.raises(ValueError, match="Exactly one"):
        create_acp_agent()

    with pytest.raises(ValueError, match="Exactly one"):
        create_acp_agent(
            agent=Agent(TestModel()),
            agent_factory=lambda session: Agent(TestModel()),
        )
