import threading
import time
import unittest

import prain.chat as chat
from tests.helpers import StubVllm, wait_for


class Chat(unittest.TestCase):
    def setUp(self):
        self.stub = StubVllm()
        self.addCleanup(self.stub.close)

    def warmer(self):
        return chat.Warmer(lambda draft: chat.send_request(
            self.stub.url, "/v1/chat/completions", chat.warm_body("m", [], draft, priority=1)))

    def test_warm_body_continues_the_half_typed_message(self):
        body = chat.warm_body("m", [{"role": "assistant", "content": "hi"}], "  Explain mix ", 1)
        self.assertEqual(body["messages"][-1], {"role": "user", "content": "Explain mix"})
        self.assertEqual(
            (body["max_tokens"], body["add_generation_prompt"], body["continue_final_message"],
             body["stream"], body["priority"]), (1, False, True, True, 1))
        self.assertNotIn("priority", chat.warm_body("m", [], "x"))

    def test_draft_body_marks_the_request_and_salts_its_cache(self):
        body = chat.draft_body("m", [{"role": "user", "content": "hi"}])
        self.assertEqual(body, {"model": "m", "messages": [{"role": "user", "content": "hi"}],
                                "stream": True, "vllm_xargs": {"prain": "draft"},
                                "cache_salt": "prain-draft"})

    def test_first_model_and_streamed_text(self):
        self.assertEqual(chat.first_model(self.stub.url), "stub-model")
        with chat.post(self.stub.url, "/v1/chat/completions",
                       chat.chat_body("stub-model", [], max_tokens=5)) as response:
            self.assertEqual(list(chat.stream_text(response)), ["Hel", "lo", "!"])
        self.assertEqual((self.stub.bodies[0]["stream"], self.stub.bodies[0]["max_tokens"]),
                         (True, 5))

    def test_warmer_keeps_one_request_in_flight_and_sends_the_newest_draft(self):
        self.stub.delay = 0.3
        warmer = self.warmer()
        warmer.update("a")
        wait_for(lambda: len(self.stub.bodies) == 1)
        warmer.update("ab")  # "a" is still in flight: these two queue up,
        warmer.update("abc")  # and only the newest is sent
        wait_for(lambda: not warmer.busy)
        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["a", "abc"])

    def test_cancel_drops_the_queue_and_aborts_the_request_in_flight(self):
        self.stub.delay = 3.0
        warmer = self.warmer()
        warmer.update("slow")
        wait_for(lambda: self.stub.bodies)
        warmer.update("queued")
        start = time.monotonic()
        warmer.cancel()
        wait_for(lambda: not warmer.busy, timeout=2.0)
        self.assertLess(time.monotonic() - start, 1.0)  # did not wait out the 3 s answer
        self.assertEqual(len(self.stub.bodies), 1)  # "queued" was never sent