from __future__ import annotations

import sys
import unittest
from pathlib import Path
from types import SimpleNamespace

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from core.loop import AgentLoop, _malformed_tool_calls  # noqa: E402


def _resp(arguments: str):
    """最小 response 桩:单条 tool_call,arguments 由测试指定。"""
    tc = SimpleNamespace(function=SimpleNamespace(name="write", arguments=arguments))
    msg = SimpleNamespace(tool_calls=[tc], content=None)
    return SimpleNamespace(choices=[SimpleNamespace(message=msg)])


GOOD = '{"path": "a.md", "content": "ok"}'
BAD = '].cells[1].merge({"path": "a.md"'  # 流式 delta 错位的典型形态


def _make_loop(stream_results, nonstream_results):
    """绕开构造器搭 AgentLoop:只装 _stream_llm 用到的属性,打桩两条取流路径。"""
    loop = object.__new__(AgentLoop)
    loop.caps = SimpleNamespace(reliable_context=64_000)
    loop.session = SimpleNamespace(messages=[], task_id="test-task")
    loop.events = []
    loop._emit = loop.events.append
    calls = {"stream": 0, "nonstream": 0}

    def collect_stream_once(llm_messages):
        calls["stream"] += 1
        return stream_results.pop(0), False

    def nonstream_once(llm_messages):
        calls["nonstream"] += 1
        return nonstream_results.pop(0)

    loop._collect_stream_once = collect_stream_once
    loop._nonstream_once = nonstream_once
    return loop, calls


class TestMalformedRetry(unittest.TestCase):
    def test_clean_stream_no_retry(self):
        loop, calls = _make_loop([_resp(GOOD)], [])
        resp, cancelled = loop._stream_llm()
        self.assertFalse(cancelled)
        self.assertEqual(_malformed_tool_calls(resp), [])
        self.assertEqual(calls, {"stream": 1, "nonstream": 0})

    def test_first_failure_falls_back_to_nonstream(self):
        """核心断言:流式只试 1 次,首败即降级非流式(不再流式重 roll)。"""
        loop, calls = _make_loop([_resp(BAD)], [_resp(GOOD)])
        resp, cancelled = loop._stream_llm()
        self.assertFalse(cancelled)
        self.assertEqual(_malformed_tool_calls(resp), [])
        self.assertEqual(calls, {"stream": 1, "nonstream": 1})
        warns = [e for e in loop.events if e.get("type") == "warn"]
        self.assertEqual(len(warns), 1)
        self.assertIn("非流式", warns[0]["msg"])

    def test_all_attempts_exhausted_returns_last(self):
        """全部尝试耗尽仍畸形 → 返回最后一次 response(交 executor 返错),不死循环。"""
        n_nonstream = AgentLoop._MAX_MALFORMED_ATTEMPTS - 1
        loop, calls = _make_loop([_resp(BAD)], [_resp(BAD)] * n_nonstream)
        resp, cancelled = loop._stream_llm()
        self.assertFalse(cancelled)
        self.assertTrue(_malformed_tool_calls(resp))
        self.assertEqual(calls, {"stream": 1, "nonstream": n_nonstream})
        warns = [e for e in loop.events if e.get("type") == "warn"]
        self.assertEqual(len(warns), AgentLoop._MAX_MALFORMED_ATTEMPTS)

    def test_cancel_mid_stream_short_circuits(self):
        loop, calls = _make_loop([], [])
        loop._collect_stream_once = lambda m: (None, True)
        resp, cancelled = loop._stream_llm()
        self.assertTrue(cancelled)
        self.assertIsNone(resp)
        self.assertEqual(calls["nonstream"], 0)


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