"""Metrics: usage forwarding from the LLM stream and first-audio-chunk (TTS TTFB) timing."""
from __future__ import annotations

import json
import os
import tempfile
import time
import unittest
from unittest import mock

os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-metrics-test-config-"))

from fastapi import FastAPI
from fastapi.testclient import TestClient

import routes.conversation as conversation


class _Resp:
    def __init__(self, status=200, lines=()):
        self.status_code, self._lines = status, lines

    def raise_for_status(self):
        pass

    def iter_lines(self):
        return iter(self._lines)


def _chunk(content=None, usage=None):
    o = {"choices": [{"delta": {"content": content}}] if content is not None else []}
    if usage is not None:
        o["usage"] = usage
    return b"data: " + json.dumps(o).encode()


class UsageParsing(unittest.TestCase):
    def test_numbers_only(self):
        self.assertEqual(conversation._usage_from_chunk({"usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}}),
                         {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12})

    def test_absent_or_junk_is_none(self):
        for o in ({}, {"usage": None}, {"usage": {}}, {"usage": {"prompt_tokens": "x", "total_tokens": True}}):
            self.assertIsNone(conversation._usage_from_chunk(o))


class RetryWithoutOptionalKeys(unittest.TestCase):
    def test_400_drops_stream_options_and_thinking(self):
        sent = []
        resps = [_Resp(400), _Resp(200)]

        def fake(u, json=None, **kw):
            sent.append(dict(json))
            return resps.pop(0)

        payload = {"stream": True, "stream_options": {"include_usage": True}, "chat_template_kwargs": {}, "think": False}
        with mock.patch.object(conversation.requests, "post", side_effect=fake):
            r = conversation._post_llm_stream("http://h/v1/chat/completions", payload, {})
        self.assertEqual(r.status_code, 200)
        self.assertIn("stream_options", sent[0])
        self.assertEqual(sent[1], {"stream": True})

    def test_no_retry_when_nothing_optional(self):
        with mock.patch.object(conversation.requests, "post", return_value=_Resp(400)) as p:
            self.assertEqual(conversation._post_llm_stream("u", {"stream": True}, {}).status_code, 400)
        self.assertEqual(p.call_count, 1)


class TurnStream(unittest.TestCase):
    def _turn(self, lines_fn, tts_delay=0.0):
        app = FastAPI()
        app.include_router(conversation.router)
        sent = []

        def fake_post(u, json=None, **kw):
            sent.append(json)
            return _Resp(lines=lines_fn())

        def fake_tts(*a, **k):
            time.sleep(tts_delay)
            return b"RIFF", "audio/wav"

        with mock.patch.object(conversation, "_load_settings", return_value={"llm_url": "http://h:17861/v1"}), \
                mock.patch.object(conversation.requests, "post", side_effect=fake_post), \
                mock.patch.object(conversation, "_preview_request_audio", side_effect=fake_tts):
            r = TestClient(app).post("/api/conversation/turn", data={"text": "Tell me about the sea please"})
        evts = [json.loads(l[5:]) for l in r.text.splitlines() if l.startswith("data:")]
        return sent, next(e for e in evts if e["type"] == "stats")

    def test_usage_forwarded_and_requested(self):
        sent, stats = self._turn(lambda: [_chunk("The sea is wide and deep. "), _chunk(usage={"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13}), b"data: [DONE]"])
        self.assertEqual(sent[0]["stream_options"], {"include_usage": True})
        self.assertEqual(stats["usage"], {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13})

    def test_no_usage_is_null_not_zero(self):
        _, stats = self._turn(lambda: [_chunk("The sea is wide and deep. "), b"data: [DONE]"])
        self.assertIsNone(stats["usage"])

    def test_tts_ms_is_first_chunk_not_llm_duration(self):
        def lines():
            yield _chunk("The sea is wide and deep. ")
            time.sleep(0.6)  # LLM keeps streaming long after the first sentence's audio is ready
            yield _chunk("It is also very blue today. ")
            yield b"data: [DONE]"
        _, stats = self._turn(lines)
        self.assertGreaterEqual(stats["llm_total_ms"], 600)
        self.assertLess(stats["tts_ms"], 400)


if __name__ == "__main__":
    unittest.main()