import json
import threading
import unittest
import urllib.error
import urllib.request

import prain.chat as chat
from prain.config import Draft, Serve
from prain.web import make_server
from tests.helpers import StubVllm, wait_for


def events(response):
    """The proxy's stream as a list: content pieces, and "<draft>"-style markers."""
    out = []
    for raw in response:
        line = raw.decode().strip()
        if line == "data: [DONE]":
            out.append("<done>")
        elif line.startswith("data: "):
            data = json.loads(line[6:])
            out.append(f"<{data['prain']}>" if "prain" in data
                       else data["choices"][0]["delta"]["content"])
    return out


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

    def start(self, **options):
        self.server = make_server(self.stub.url, 0, **options)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        self.addCleanup(self.server.server_close)
        self.addCleanup(self.server.shutdown)
        self.url = f"http://127.0.0.1:{self.server.server_port}"

    def warm(self, session, draft):
        chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read()

    def ask(self, session, text):
        return chat.post(self.url, "/chat", {"session": session,
                                             "messages": [{"role": "user", "content": text}]})

    def test_page_warm_up_and_chat(self):
        with urllib.request.urlopen(self.url) as response:
            self.assertIn(b'<textarea id="draft"', response.read())
        self.warm("s1", "Why is ")
        wait_for(lambda: self.stub.bodies)
        self.assertEqual(self.stub.bodies[0]["messages"], [{"role": "user", "content": "Why is"}])
        self.assertEqual(self.stub.bodies[0]["priority"], 1)
        with self.ask("s1", "Why is") as r:
            self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
        self.assertEqual(self.stub.bodies[1]["model"], "stub-model")
        self.assertNotIn("vllm_xargs", self.stub.bodies[1])

    def test_sessions_warm_up_independently(self):
        self.stub.delay = 0.3
        self.warm("alice", "first")
        self.warm("bob", "second")  # alice's warm-up is in flight; bob's is not dropped
        wait_for(lambda: len(self.stub.bodies) == 2)
        self.assertEqual({b["messages"][-1]["content"] for b in self.stub.bodies},
                         {"first", "second"})
        self.assertEqual(set(self.server.sessions), {"alice", "bob"})

    def test_no_warm_up_while_another_answer_streams(self):
        self.server.streams[0] = 1  # someone else's answer is streaming
        self.warm("bob", "typing along")
        self.server.streams[0] = 0
        self.warm("bob", "now")
        wait_for(lambda: self.stub.bodies)
        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["now"])

    def test_always_mode_warms_up_while_a_request_slot_is_free(self):
        self.start(warmup="always", serve=Serve(max_requests=3))
        self.server.streams[0] = 3  # every vLLM request slot is busy
        self.warm("bob", "no slot")
        self.server.streams[0] = 2
        self.warm("bob", "one slot")
        wait_for(lambda: self.stub.bodies)
        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["one slot"])

    def test_vllm_errors_reach_the_browser(self):
        self.stub.fail = 400
        with self.assertRaises(urllib.error.HTTPError) as caught:
            chat.post(self.url, "/chat", {"session": "s", "messages": []})
        self.assertEqual(caught.exception.code, 400)
        with caught.exception as error:
            self.assertIn(b"rejected", error.read())
        wait_for(lambda: self.server.streams[0] == 0)

    def test_second_answer_drafts_now_and_refines_when_the_first_one_ends(self):
        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5  # full answers are slow
        alice = []
        thread = threading.Thread(target=lambda: alice.extend(events(self.ask("alice", "a"))))
        thread.start()
        wait_for(lambda: self.server.full[0] == 1)
        with self.ask("bob", "b") as r:
            bob = events(r)
        thread.join()
        self.assertEqual(alice, ["Hel", "lo", "!", "<done>"])
        self.assertEqual(bob, ["<draft>", "Hel", "lo", "!", "<refine>", "Hel", "lo", "!", "<done>"])
        self.assertEqual([(b["messages"][0]["content"], b.get("vllm_xargs"), b.get("cache_salt"))
                          for b in self.stub.bodies],
                         [("a", None, None), ("b", {"prain": "draft"}, "prain-draft"),
                          ("b", None, None)])
        wait_for(lambda: self.server.full[0] == 0 and self.server.streams[0] == 0)

    def test_a_running_draft_is_cut_short_once_a_full_slot_frees(self):
        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.3
        self.stub.pause = lambda body: 0.6 if "vllm_xargs" in body else 0.0  # drafts trickle
        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
        thread.start()
        wait_for(lambda: self.server.full[0] == 1)
        with self.ask("bob", "b") as r:  # alice ends at 0.3 s, between bob's "Hel" and "lo"
            self.assertEqual(events(r), ["<draft>", "Hel", "lo", "<refine>", "Hel", "lo", "!",
                                         "<done>"])
        thread.join()

    def test_refine_none_keeps_the_draft_and_more_full_slots_mean_no_draft(self):
        self.start(draft=Draft(refine="none"))
        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5
        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
        thread.start()
        wait_for(lambda: self.server.full[0] == 1)
        with self.ask("bob", "b") as r:
            self.assertEqual(events(r), ["<draft>", "Hel", "lo", "!", "<done>"])
        thread.join()
        self.start(serve=Serve(max_requests=2), draft=Draft(full_requests=2))
        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
        thread.start()
        wait_for(lambda: self.server.full[0] == 1)
        with self.ask("bob", "b") as r:
            self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
        thread.join()
        self.assertEqual(sum("vllm_xargs" in b for b in self.stub.bodies), 1)