"""Pausing a reply to ask the reader something.

Driven through the real generation loop with a scripted endpoint, because the
things worth pinning here are all about the loop: that Stop still works while
nothing is streaming, that a decision lands at the right index, and that the
card clears itself.
"""

from __future__ import annotations

import asyncio

import pytest
from sqlalchemy import select

from lembas.db.models import ROLE_ASSISTANT, Chat, Connection, Message, Model
from lembas.services import generation as generation_service
from lembas.services import interaction, settings_store


def _chat_that_can_ask(db, user_id):
    connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
    db.commit()
    chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
    db.add(chat)
    db.commit()
    db.add(Message(chat_id=chat.id, role="user", content="Which one?", complete=True))
    db.commit()
    assistant = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", complete=False)
    db.add(assistant)
    db.commit()
    return chat.id, assistant.id


def _ask_chunk(question: str, options: list[str] | None = None, *, index: int = 0, call_id="c1"):
    import json as _json

    arguments = {"question": question}
    if options is not None:
        arguments["options"] = options
    return {
        "choices": [
            {
                "delta": {
                    "tool_calls": [
                        {
                            "index": index,
                            "id": call_id,
                            "function": {
                                "name": "ask_user",
                                "arguments": _json.dumps(arguments),
                            },
                        }
                    ]
                }
            }
        ]
    }


def _text_chunk(text: str) -> dict:
    return {"choices": [{"delta": {"content": text}}]}


def _stub_stream(rounds, seen):
    async def stream_chat(_endpoint, payload):
        seen.append(payload)
        for chunk in rounds[min(len(seen) - 1, len(rounds) - 1)]:
            yield chunk

    return stream_chat


async def _until_paused(generation, *, timeout: float = 2.0):
    """Wait for the card to go up."""
    deadline = asyncio.get_running_loop().time() + timeout
    while asyncio.get_running_loop().time() < deadline:
        if generation.pending is not None:
            return generation.pending
        await asyncio.sleep(0.01)
    raise AssertionError("the reply never paused")


@pytest.fixture
def scripted(db, user_id, monkeypatch):
    """A reply that asks one question, then answers with whatever it was told."""
    chat_id, message_id = _chat_that_can_ask(db, user_id)
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [[_ask_chunk("Tea or coffee?", ["Tea", "Coffee"])], [_text_chunk("Right you are.")]],
            payloads,
        ),
    )
    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    return generation, payloads, chat_id, message_id


# --- The tool is offered at all ----------------------------------------------
def test_ask_user_is_offered_to_an_ordinary_chat(db, user_id):
    """Not an agent feature. A model in a plain conversation should be able to
    stop and ask which of two things you meant."""
    from lembas.db.models import User
    from lembas.services import tools as tools_service

    chat_id, _message_id = _chat_that_can_ask(db, user_id)
    chat = db.get(Chat, chat_id)
    offered = tools_service.resolve_tools(db, chat, db.get(User, user_id))
    assert "ask_user" in offered.by_name


def test_ask_user_is_withheld_without_the_permission(db, user_id):
    from lembas.db.models import User
    from lembas.services import tools as tools_service

    user = db.get(User, user_id)
    # Administrators are given every permission, so the baseline only bites a
    # plain account.
    user.role = "user"
    settings_store.update(db, {"default_permissions": {"tools.ask": False}})
    db.commit()

    chat = db.get(Chat, _chat_that_can_ask(db, user_id)[0])
    assert "ask_user" not in tools_service.resolve_tools(db, chat, user).by_name


# --- The pause ---------------------------------------------------------------
async def test_the_reply_pauses_and_the_card_describes_the_question(scripted):
    generation, _payloads, _chat_id, _message_id = scripted
    task = asyncio.create_task(generation_service._run(generation))

    pending = await _until_paused(generation)
    assert pending.kind == interaction.KIND_QUESTION
    assert pending.items[0].title == "Tea or coffee?"
    assert pending.options == ("Tea", "Coffee")
    assert "Waiting for your answer" in generation.status

    pending.resolve(interaction.ANSWER, text="Tea")
    await task


async def test_the_answer_reaches_the_model_as_a_tool_result(scripted):
    generation, payloads, _chat_id, _message_id = scripted
    task = asyncio.create_task(generation_service._run(generation))

    pending = await _until_paused(generation)
    pending.resolve(interaction.ANSWER, text="Coffee, please")
    await task

    turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
    assert len(turns) == 1
    assert "Coffee, please" in turns[0]["content"]
    assert turns[0]["tool_call_id"] == "c1"


async def test_the_card_is_cleared_once_it_is_answered(scripted):
    generation, _payloads, _chat_id, _message_id = scripted
    task = asyncio.create_task(generation_service._run(generation))

    pending = await _until_paused(generation)
    version = generation.version
    pending.resolve(interaction.ANSWER, text="Tea")
    await task

    assert generation.pending is None
    assert generation.version > version, "clearing has to bump the version or no frame is sent"


async def test_the_transcript_keeps_what_was_asked_and_answered(scripted):
    generation, _payloads, _chat_id, _message_id = scripted
    task = asyncio.create_task(generation_service._run(generation))

    pending = await _until_paused(generation)
    pending.resolve(interaction.ANSWER, text="Tea")
    await task

    event = generation.tool_events[0]
    assert event["kind"] == "ask"
    assert event["query"] == "Tea or coffee?"
    assert event["text"] == "Tea"


# --- Stop, while nothing is streaming ----------------------------------------
async def test_stop_ends_a_reply_that_is_waiting_for_an_answer(db, user_id, monkeypatch):
    """`cancel` is read only between streamed chunks, and there are no chunks
    while the card is up. Without the wakeup in request_stop the button does
    nothing at all here."""
    chat_id, message_id = _chat_that_can_ask(db, user_id)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream([[_ask_chunk("Tea or coffee?")], [_text_chunk("unreachable")]], []),
    )

    generation = generation_service.ensure(chat_id, message_id)
    await _until_paused(generation)

    assert generation_service.request_stop(message_id) is True
    await asyncio.wait_for(asyncio.shield(generation_service._TASKS[message_id]), timeout=2)

    assert generation.stopped is True
    assert generation.pending is None


# --- Timeout ------------------------------------------------------------------
async def test_an_unanswered_question_expires_and_the_reply_finishes(db, user_id, monkeypatch):
    chat_id, message_id = _chat_that_can_ask(db, user_id)
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream([[_ask_chunk("Tea or coffee?")], [_text_chunk("Never mind.")]], payloads),
    )
    # The clamp floor is 60s, so the timeout is forced directly rather than
    # through the settings.
    monkeypatch.setattr(
        generation_service.tools_service,
        "context_for",
        lambda *a, **k: _fast_context(*a, **k),
    )

    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    await asyncio.wait_for(generation_service._run(generation), timeout=5)

    turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
    assert "did not answer" in turns[0]["content"]
    assert generation.tool_events[0]["status"] == "error"


def _fast_context(db, user, chat=None, *, tools=None):
    from lembas.services import tools as tools_service

    context = tools_service.ToolContext(
        owner_id=user.id if user else "",
        tools=tools.by_name if tools is not None else None,
    )
    context.interaction_timeout = 0.2
    return context


# --- The primitive on its own -------------------------------------------------
def test_resolving_twice_only_counts_once():
    """Two tabs, or a double click. The second answer must not win."""

    async def go():
        pause = interaction.build("abc", [_item()], timeout=5)
        assert pause.resolve(interaction.ANSWER, text="first") is True
        assert pause.resolve(interaction.ANSWER, text="second") is False
        assert (await pause._future).text == "first"

    asyncio.run(go())


def test_an_interruption_with_no_future_cannot_be_resolved():
    pause = interaction.Interruption(id="abc", items=(_item(),))
    assert pause.resolve(interaction.ANSWER) is False


def _item() -> interaction.Item:
    return interaction.Item(
        index=0,
        kind=interaction.KIND_QUESTION,
        tool_name="ask_user",
        title="Tea or coffee?",
    )


# --- Resolving one -------------------------------------------------------------
# `answer()` is exercised in-process rather than over the TestClient, because a
# future belongs to the loop that made it and TestClient runs the app on its
# own. In production both are the single uvicorn loop, which is the arrangement
# the in-process registry already requires.
async def test_answer_finds_the_pause_and_resolves_it(db, user_id):
    chat_id, message_id = _chat_that_can_ask(db, user_id)
    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    generation.pending = interaction.build("pause-1", [_item()], timeout=30)
    generation_service._RUNNING[message_id] = generation
    try:
        assert generation_service.answer(chat_id, "pause-1", choice="Tea", text="") is True
        assert (await generation.pending._future).text == "Tea"
    finally:
        generation_service._RUNNING.pop(message_id, None)


async def test_answer_ignores_a_pause_in_another_chat(db, user_id):
    """The endpoint checks ownership of the chat, so the lookup is scoped to it
    -- an id on its own would not be an authorisation."""
    chat_id, message_id = _chat_that_can_ask(db, user_id)
    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    generation.pending = interaction.build("pause-1", [_item()], timeout=30)
    generation_service._RUNNING[message_id] = generation
    try:
        assert generation_service.answer("another-chat", "pause-1", choice="x", text="") is False
        assert not generation.pending._future.done()
    finally:
        generation_service._RUNNING.pop(message_id, None)


async def test_an_allow_choice_is_kept_as_a_verdict_not_as_typed_text():
    """"allow" is a decision, not something somebody wrote in the box."""
    pause = interaction.build("p", [_item()], timeout=5)
    pause.resolve(interaction.ALLOW)
    reply = await pause._future
    assert reply.permitted is True
    assert reply.outcome == interaction.ALLOW


def test_answering_something_that_has_gone_says_so(client, db, registered, user_id):
    chat_id, _message_id = _chat_that_can_ask(db, user_id)
    response = client.post(f"/api/chats/{chat_id}/interaction/nope", data={"choice": "Tea"})
    assert response.status_code == 204
    assert "no longer waiting" in response.headers["HX-Trigger"]


def test_another_account_cannot_answer_your_question(client, db, registered, user_id):
    """Without the ownership check, a guessed id would be answering -- and
    later, approving a command in -- somebody else's conversation."""
    from lembas.db.models import User

    chat_id, message_id = _chat_that_can_ask(db, user_id)
    generation = generation_service.Generation(chat_id=chat_id, message_id=message_id)
    # No future: the request must be refused before anything tries to resolve
    # it, so there is nothing here for it to reach.
    pause = interaction.Interruption(id="pause-2", items=(_item(),))
    generation.pending = pause
    generation_service._RUNNING[message_id] = generation
    try:
        client.post("/auth/logout")
        client.post(
            "/auth/register",
            data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"},
            follow_redirects=False,
        )
        intruder = db.scalar(select(User).where(User.email == "sam@shire.test"))
        intruder.role = "user"
        intruder.active = True
        db.commit()
        client.post(
            "/auth/login",
            data={"email": "sam@shire.test", "password": "correct horse battery"},
            follow_redirects=False,
        )

        response = client.post(
            f"/api/chats/{chat_id}/interaction/pause-2", data={"choice": "Tea"}
        )
        assert response.status_code == 404
        assert pause is generation.pending, "still waiting for the person it belongs to"
    finally:
        generation_service._RUNNING.pop(message_id, None)


# --- The card ------------------------------------------------------------------
def _render(pending) -> str:
    from lembas.api.chats import _ask_html

    return _ask_html("chat-1", pending)


def test_nothing_pending_renders_nothing():
    """The frame is sent unconditionally so the card can clear itself. An empty
    string is how it does that."""
    assert _render(None) == ""


def test_the_card_shows_the_question_and_its_options():
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                kind=interaction.KIND_QUESTION,
                tool_name="ask_user",
                title="Tea or coffee?",
                options=("Tea", "Coffee"),
            ),
        ),
    )
    html = _render(pause)
    assert "Tea or coffee?" in html
    assert 'value="Tea"' in html and 'value="Coffee"' in html
    assert 'hx-post="/api/chats/chat-1/interaction/p1"' in html
    assert "The model is asking you" in html, "attributed to the model, not to LLeMbas"


def test_the_card_never_offers_a_password_field():
    """A model talked into asking for a credential must not be handed a field
    that looks built for one."""
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                kind=interaction.KIND_QUESTION,
                tool_name="ask_user",
                title="Confirm your password to continue:",
            ),
        ),
    )
    html = _render(pause)
    assert 'type="password"' not in html
    assert 'type="text"' in html


def test_everything_on_the_card_is_escaped():
    """The question is model output, and the model may have been reading
    somebody else's file a moment ago."""
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                kind=interaction.KIND_QUESTION,
                tool_name="ask_user",
                title="<img src=x onerror=alert(1)>",
                detail="rm -rf / <script>",
                options=('"><script>bad</script>',),
            ),
        ),
    )
    html = _render(pause)
    assert "<img" not in html
    assert "<script>" not in html
    assert "&lt;img" in html