from __future__ import annotations

import sys
import threading
import time
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, family="test", variant="t")
    loop.session = SimpleNamespace(messages=[], task_id="test-task")
    loop.user_id = "test-user"  # 无 DB 环境:_log_malformed_args 的落库路径应静默跳过
    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"])
        # 每个 attempt 重发 llm_start:warn 关掉前端占位段后,靠它重建「思考中」跳秒
        starts = [e for e in loop.events if e.get("type") == "llm_start"]
        self.assertEqual(len(starts), 2)

    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)

    def test_cancel_during_nonstream_retry(self):
        """非流式重试期间 cancel(_nonstream_once 返回 None)→ (None, True),不再续试。"""
        loop, calls = _make_loop([_resp(BAD)], [None, _resp(GOOD)])
        resp, cancelled = loop._stream_llm()
        self.assertTrue(cancelled)
        self.assertIsNone(resp)
        self.assertEqual(calls, {"stream": 1, "nonstream": 1})

    def test_nonstream_once_polls_cancel(self):
        """真 _nonstream_once:llm.chat 阻塞时点停止,应在 poll 拍内返回 None,
        不等调用整个返回(此前同步阻塞,分钟级生成期间停止按钮无效)。"""
        loop = object.__new__(AgentLoop)
        loop.caps = SimpleNamespace(default_reasoning_effort=None)
        loop.executor = SimpleNamespace(schemas=lambda: [])
        loop.cancel_check = lambda: True
        block = threading.Event()

        def slow_chat(**kwargs):
            block.wait(30)
            return _resp(GOOD)

        loop.llm = SimpleNamespace(chat=slow_chat)
        t0 = time.monotonic()
        resp = loop._nonstream_once([])
        elapsed = time.monotonic() - t0
        block.set()  # 释放弃养线程,别让它拖住测试进程退出
        self.assertIsNone(resp)
        self.assertLess(elapsed, 5)

    def test_nonstream_once_returns_response_when_not_cancelled(self):
        """不取消时行为与原同步版等价:拿到 response 并补 emit 整段 text。"""
        loop = object.__new__(AgentLoop)
        loop.caps = SimpleNamespace(default_reasoning_effort=None)
        loop.executor = SimpleNamespace(schemas=lambda: [])
        loop.cancel_check = None
        loop.events = []
        loop._emit = loop.events.append
        good = _resp(GOOD)
        good.choices[0].message.content = "hello"
        loop.llm = SimpleNamespace(chat=lambda **kwargs: good)
        resp = loop._nonstream_once([])
        self.assertIs(resp, good)
        texts = [e for e in loop.events if e.get("type") == "text"]
        self.assertEqual(len(texts), 1)
        self.assertEqual(texts[0]["delta"], "hello")


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