"""Chat finish trailers retain usage without admitting late response data, and
replayed assistant messages keep the reasoning field the provider speaks."""

import json
import unittest

from harness import Albedo, Provider, Reply, python, text


def chunk(delta, finish=None):
    return {
        "id": "fixture",
        "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
    }


class ChatStreamTest(unittest.TestCase):
    def setUp(self):
        self.provider = Provider(lambda _: Reply("text"))
        self.addCleanup(self.provider.close)

    def request_row(self, app, session):
        with app.api(f"/sessions/{session}/context?view=requests") as response:
            [row] = json.load(response)["items"]
        return row

    def test_repeated_finish_retains_trailing_usage_and_commits_answer_once(self):
        # OpenRouter's shape: both finish chunks repeat the assistant role with
        # empty content, and the second carries usage.
        finish = chunk({"content": "", "role": "assistant", "reasoning": None}, "stop")
        trailer = chunk({"content": "", "role": "assistant"}, "stop")
        trailer["usage"] = {"prompt_tokens": 123, "completion_tokens": 7}
        self.provider.script = lambda _: Reply(
            "text",
            events=[chunk({"content": "answer"}), finish, trailer],
        )
        with Albedo(self.provider, protocol="chat_completions") as app:
            session = app.session()
            app.prompt(session, "answer once").close()
            app.idle(session)
            row = self.request_row(app, session)
            self.assertEqual(row["outcome"], "completed")
            self.assertEqual(row["prompt_tokens"]["observed"], 123)
            self.assertEqual(row["completion_tokens"]["observed"], 7)
            [answer] = [
                entry
                for entry in app.history(session)["items"]
                if entry["kind"] == "assistant"
            ]
            self.assertEqual(answer["content"], [{"kind": "text", "text": "answer"}])

    def test_post_finish_choices_must_be_empty_and_repeat_the_same_finish(self):
        with Albedo(self.provider, protocol="chat_completions") as app:
            for delta, finish in [
                ({}, "length"),
                ({}, None),
                ({"content": "late"}, "stop"),
                ({"refusal": "late"}, "stop"),
                ({"reasoning": "late"}, "stop"),
                ({"reasoning_content": "late"}, "stop"),
                (
                    {"reasoning_details": [{"type": "reasoning.text", "index": 0}]},
                    "stop",
                ),
                ({"tool_calls": [{"index": 0}]}, "stop"),
            ]:
                with self.subTest(delta=delta, finish=finish):
                    self.provider.script = lambda _: Reply(
                        "text",
                        events=[
                            chunk({"content": "answer"}, "stop"),
                            chunk(delta, finish),
                        ],
                    )
                    session = app.session()
                    app.prompt(session, "reject the late choice").close()
                    app.idle(session)
                    row = self.request_row(app, session)
                    self.assertEqual(row["outcome"], "failed")
                    self.assertIn(
                        "chat choice after finish reason", row["failure"]["detail"]
                    )
                    # What streamed before the late choice is kept as
                    # unfinished output; the late choice itself is not.
                    answers = [
                        entry["content"]
                        for entry in app.history(session)["items"]
                        if entry["kind"] == "assistant"
                    ]
                    self.assertEqual(answers, [[{"kind": "text", "text": "answer"}]])

    def test_every_assistant_message_echoes_reasoning_once_the_provider_speaks_it(self):
        # DeepSeek's thinking mode rejects a tools request whose assistant
        # messages lack reasoning_content, even one that never reasoned.
        self.provider.script = lambda _: [
            python("x = 1", reasoning="weigh it"),
            python("x = 2"),
            text("done"),
        ]
        with Albedo(self.provider, protocol="chat_completions") as app:
            session = app.session()
            app.prompt(session, "two steps").close()
            app.idle(session)
            last = self.provider.requests[-1]["request"]["messages"]
            echoed = [
                message.get("reasoning_content")
                for message in last
                if message["role"] == "assistant"
            ]
            self.assertEqual(echoed, ["weigh it", ""])