"""python -m unittest discover -s tests -p test_stream_tool_calls.py -v"""
import json
import unittest
from pathlib import Path
from unittest.mock import patch

import torch
from fastapi.testclient import TestClient
from openai import OpenAI
from transformers import AutoTokenizer

from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
from scripts import serve_openai_api as server


TOOL = '<tool_call>{"name":"weather","arguments":{"city":"杭州"}}</tool_call>'
TOOLS = [{'type': 'function', 'function': {'name': 'weather', 'parameters': {
    'type': 'object', 'properties': {'city': {'type': 'string'}}, 'required': ['city']}}}]


class ScriptedModel:
    def __init__(self, chunks, tokenize=False):
        self.chunks, self.tokenize = chunks, tokenize

    def generate(self, input_ids, streamer, **kwargs):
        if self.tokenize:
            streamer.put(input_ids.cpu())
            for token in streamer.tokenizer.encode(''.join(self.chunks), add_special_tokens=False):
                streamer.put(torch.tensor([token]))
            streamer.end()
        else:
            for chunk in self.chunks:
                streamer.on_finalized_text(chunk)
            streamer.on_finalized_text('', stream_end=True)


class StreamToolTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tokenizer = AutoTokenizer.from_pretrained(
            Path(__file__).resolve().parents[1] / 'model', local_files_only=True)

    def setUp(self):
        globals_patch = patch.multiple(server, tokenizer=self.tokenizer, device='cpu',
                                       model=ScriptedModel([]), create=True)
        globals_patch.start()
        self.addCleanup(globals_patch.stop)
        self.http = TestClient(server.app)
        self.addCleanup(self.http.close)
        self.request = {'model': 'minimind', 'messages': [{'role': 'user', 'content': 'Hi'}]}

    def frames(self, chunks, thinking=False):
        server.model = ScriptedModel(chunks)
        response = self.http.post('/v1/chat/completions', json={
            **self.request, 'stream': True, 'open_thinking': thinking})
        self.assertEqual(response.status_code, 200)
        self.assertIn('text/event-stream', response.headers['content-type'])
        return [json.loads(line[6:]) for line in response.text.splitlines()
                if line.startswith('data: ')]

    def assert_output(self, frames, content, count, reasoning=''):
        self.assertFalse(any('error' in frame for frame in frames), frames)
        deltas = [frame['choices'][0]['delta'] for frame in frames]
        self.assertEqual(''.join(d.get('content', '') for d in deltas), content)
        self.assertEqual(''.join(d.get('reasoning_content', '') for d in deltas), reasoning)
        calls = [call for d in deltas for call in d.get('tool_calls', [])]
        self.assertEqual(len(calls), count)
        self.assertEqual([call['index'] for call in calls], list(range(count)))
        self.assertEqual(len({call['id'] for call in calls}), count)
        for call in calls:
            self.assertEqual(call['function']['name'], 'weather')
            self.assertEqual(json.loads(call['function']['arguments']), {'city': '杭州'})
        self.assertEqual(frames[-1]['choices'][0]['finish_reason'], 'tool_calls' if count else 'stop')

    def test_split_markers_and_multiple_calls(self):
        text = 'Before ' + TOOL + ' between ' + TOOL + ' after'
        for chunks in ([text], list(text), ['Before <tool_', text[len('Before <tool_'):]]):
            with self.subTest(chunks=len(chunks)):
                self.assert_output(self.frames(chunks), 'Before  between  after', 2)

    def test_thinking_boundary_and_tool_text_inside_reasoning(self):
        reasoning = 'Consider ' + TOOL + '\n'
        text = reasoning + '</think>\n\n' + TOOL + ' answer'
        self.assert_output(self.frames(list(text), thinking=True), ' answer', 1, reasoning)

    def test_invalid_and_unfinished_calls_preserve_text(self):
        for text in ('plain < text', 'end <tool_', '<tool_call>{"name":',
                     '<tool_call>bad json</tool_call>', '<tool_call>[]</tool_call>'):
            with self.subTest(text=text):
                self.assert_output(self.frames(list(text)), text, 0)
        invalid = '<tool_call>bad json</tool_call>'
        self.assert_output(self.frames([invalid, TOOL]), invalid, 1)

    def test_empty_and_unfinished_thinking(self):
        self.assert_output(self.frames([]), '', 0)
        self.assert_output(self.frames(['reason </thi'], thinking=True), '', 0, 'reason </thi')

    def test_plain_content_and_tool_call_arrive_before_generation_ends(self):
        def generate(input_ids, streamer, **kwargs):
            streamer.on_finalized_text('first ')
            self.assertTrue(continue_generation.wait(5))
            streamer.on_finalized_text(TOOL)
            self.assertTrue(finish_generation.wait(5))
            streamer.on_finalized_text(' last', stream_end=True)

        from threading import Event
        continue_generation, finish_generation = Event(), Event()
        server.model.generate = generate
        stream = server.generate_stream_response(self.request['messages'], 0.7, 0.92, 128)
        self.addCleanup(continue_generation.set)
        self.addCleanup(finish_generation.set)
        first = json.loads(next(stream))
        if first['choices'][0]['delta'].get('role') == 'assistant':
            first = json.loads(next(stream))
        self.assertEqual(first['choices'][0]['delta']['content'], 'first ')
        continue_generation.set()
        call = json.loads(next(stream))
        self.assertEqual(call['choices'][0]['delta']['tool_calls'][0]['function']['name'], 'weather')
        finish_generation.set()
        self.assert_output([first, call] + [json.loads(frame) for frame in stream], 'first  last', 1)

    def test_real_text_streamer_and_openai_client(self):
        server.model = ScriptedModel(['Before ', TOOL, ' after'], tokenize=True)
        with OpenAI(api_key='local-test', base_url='http://testserver/v1',
                    http_client=self.http, max_retries=0) as client:
            chunks = list(client.chat.completions.create(**self.request, tools=TOOLS, stream=True))
        self.assertEqual(''.join(c.choices[0].delta.content or '' for c in chunks), 'Before  after')
        calls = [call for c in chunks for call in c.choices[0].delta.tool_calls or []]
        self.assertEqual(len(calls), 1)
        self.assertEqual(calls[0].index, 0)
        self.assertEqual(json.loads(calls[0].function.arguments), {'city': '杭州'})

    def test_tiny_minimind_generation(self):
        torch.manual_seed(42)
        server.model = MiniMindForCausalLM(MiniMindConfig(
            hidden_size=32, num_hidden_layers=1, num_attention_heads=4,
            num_key_value_heads=2, intermediate_size=64,
            vocab_size=len(self.tokenizer), max_position_embeddings=256)).eval()
        response = self.http.post('/v1/chat/completions', json={
            **self.request, 'stream': True, 'max_tokens': 4})
        frames = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith('data: ')]
        self.assertFalse(any('error' in frame for frame in frames), frames)
        self.assertEqual(frames[-1]['choices'][0]['finish_reason'], 'stop')


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