from __future__ import annotations

import json
import unittest
from pathlib import Path
from unittest.mock import patch

from backend.features.mentor.agent import MentorSkill, chat_with_mentor, stream_with_mentor


class FakeStreamResponse:
    def __init__(self, lines: list[bytes]) -> None:
        self.lines = lines

    def __enter__(self):
        return iter(self.lines)

    def __exit__(self, exc_type, exc_value, traceback):
        return False


class MentorStreamTests(unittest.TestCase):
    def setUp(self) -> None:
        self.skill = MentorSkill(
            skill_id="test-mentor",
            name="测试老师",
            description="测试",
            tagline="先看事实",
            focus=("纪律",),
            content="只做条件化判断。",
            path=Path("SKILL.md"),
        )
        self.lines = [
            b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
            b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
            b"data: [DONE]\n",
        ]

    def test_stream_requests_upstream_streaming_and_yields_deltas(self):
        captured = {}

        def open_request(request, timeout):
            captured["payload"] = json.loads(request.data.decode("utf-8"))
            captured["accept"] = request.headers.get("Accept")
            return FakeStreamResponse(self.lines)

        with patch("backend.llm.transport.urllib.request.urlopen", side_effect=open_request):
            chunks = list(
                stream_with_mentor(
                    self.skill, {"data_trade_date": "20260723"}, "怎么看？", [],
                    "key", "https://example.test/v1", "model",
                )
            )

        self.assertEqual(chunks, ["first", " second"])
        self.assertTrue(captured["payload"]["stream"])
        self.assertEqual(captured["accept"], "text/event-stream")

    def test_non_streaming_compatibility_wrapper_collects_chunks(self):
        with patch(
            "backend.llm.transport.urllib.request.urlopen",
            return_value=FakeStreamResponse(self.lines),
        ):
            result = chat_with_mentor(
                self.skill, {}, "怎么看？", [], "key", "https://example.test/v1", "model"
            )
        self.assertEqual(result["answer"], "first second")

    def test_final_full_message_does_not_duplicate_streamed_deltas(self):
        lines = [
            b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
            b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
            b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
            b"data: [DONE]\n",
        ]
        with patch(
            "backend.llm.transport.urllib.request.urlopen",
            return_value=FakeStreamResponse(lines),
        ):
            chunks = list(
                stream_with_mentor(
                    self.skill, {}, "question", [], "key",
                    "https://example.test/v1", "model",
                )
            )
        self.assertEqual(chunks, ["first", " second"])

    def test_snapshot_only_stream_emits_only_new_suffix(self):
        lines = [
            b'data: {"choices":[{"message":{"content":"first"}}]}\n',
            b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
            b"data: [DONE]\n",
        ]
        with patch(
            "backend.llm.transport.urllib.request.urlopen",
            return_value=FakeStreamResponse(lines),
        ):
            chunks = list(
                stream_with_mentor(
                    self.skill, {}, "question", [], "key",
                    "https://example.test/v1", "model",
                )
            )
        self.assertEqual(chunks, ["first", " second"])


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