import fcntl
import os
import pty
import struct
import subprocess
import sys
import termios
import threading
import unittest

from prain.tui import Chat
from tests.helpers import StubVllm, wait_for


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

    def test_typing_warms_up_and_enter_streams_the_answer(self):
        chat = Chat(self.stub.url, "stub-model")
        chat.type("Why is the sky")
        wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
        self.assertEqual(self.stub.bodies[0]["priority"], 1)
        chat.send()
        chat.type("next question")  # an answer is streaming: no warm-up in "idle" mode
        wait_for(lambda: chat.poll() or chat.reply is None)
        self.assertEqual(chat.messages, [{"role": "user", "content": "Why is the sky"},
                                         {"role": "assistant", "content": "Hello!"}])
        self.assertEqual(len(self.stub.bodies), 2)
        self.assertIn("tok/s", chat.status)

    def test_warmup_off_and_cancel_on_send(self):
        off = Chat(self.stub.url, "stub-model", warmup="off")
        off.type("hello")
        self.assertFalse(off.warmer.busy)
        self.stub.delay = 3.0
        chat = Chat(self.stub.url, "stub-model")
        chat.type("a slow warm-up")
        wait_for(lambda: self.stub.bodies)
        chat.send()  # aborts the warm-up instead of queueing behind it
        wait_for(lambda: not chat.warmer.busy, timeout=2.0)

    def test_curses_screen_runs_in_a_real_terminal(self):
        master, slave = pty.openpty()
        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 80, 0, 0))
        tui = subprocess.Popen([sys.executable, "-m", "prain.cli", "tui", "--url", self.stub.url],
                               stdin=slave, stdout=slave, stderr=slave,
                               env={**os.environ, "TERM": "xterm"}, start_new_session=True)
        os.close(slave)
        screen = bytearray()
        threading.Thread(target=self._drain, args=(master, screen), daemon=True).start()
        wait_for(lambda: b"stub-model" in screen)  # status line drawn
        os.write(master, b"hi\r")
        wait_for(lambda: any(body.get("stream") and "priority" not in body
                             for body in self.stub.bodies))
        wait_for(lambda: b"tok/s" in screen)
        os.write(master, b"\x1b")
        self.assertEqual(tui.wait(timeout=5), 0)

    @staticmethod
    def _drain(fd, screen):
        while True:
            try:
                data = os.read(fd, 4096)
            except OSError:
                return
            if not data:
                return
            screen.extend(data)