"""§8.8 Phase 2 折叠逻辑单测:切点选择 / 内存-prepared 映射 / Session.apply_fold /
maybe_fold 编排(DB 与 LLM 全打桩,不碰真实 PG)。"""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from uuid import uuid4

import core.context_fold as cf
from core.context_fold import (
    FOLD_TAIL_RATIO,
    MIN_FOLD_MESSAGES,
    SUMMARY_NOTE_HEADER,
    _kth_user_index,
    _user_ordinal,
    build_summary_note,
    find_cutoff,
)
from core.session import Session


def _mk_window(n_turns: int, content_chars: int = 200) -> list:
    """构造 n_turns 轮 user/assistant 对话(等长消息,方便按条数推算体量占比)。"""
    msgs = []
    for i in range(n_turns):
        msgs.append({"role": "user", "content": f"u{i:03d}" + "x" * content_chars})
        msgs.append({"role": "assistant", "content": f"a{i:03d}" + "y" * content_chars})
    return msgs


class FindCutoffTests(unittest.TestCase):
    def test_cutoff_is_user_message_and_tail_ratio_respected(self) -> None:
        messages = [{"role": "system", "content": "rules"}] + _mk_window(30)
        cutoff = find_cutoff(messages, head_len=1)
        self.assertIsNotNone(cutoff)
        self.assertEqual(messages[cutoff].get("role"), "user")
        # 尾部(切点起)体量 >= FOLD_TAIL_RATIO(等长消息按条数近似)
        tail = len(messages) - cutoff
        window = len(messages) - 1
        self.assertGreaterEqual(tail / window, FOLD_TAIL_RATIO * 0.9)
        # 可折叠段不为空且够长
        self.assertGreaterEqual(cutoff - 1, MIN_FOLD_MESSAGES)

    def test_too_few_messages_returns_none(self) -> None:
        messages = [{"role": "system", "content": "rules"}] + _mk_window(3)
        self.assertIsNone(find_cutoff(messages, head_len=1))

    def test_empty_window_returns_none(self) -> None:
        self.assertIsNone(find_cutoff([{"role": "system", "content": "s"}], head_len=1))

    def test_cutoff_never_splits_tool_pairing(self) -> None:
        # assistant.tool_calls 后面跟 tool 结果 —— 切点是 user 消息,天然不落在中间
        messages = [{"role": "system", "content": "rules"}]
        for i in range(20):
            messages.append({"role": "user", "content": f"u{i}" + "x" * 200})
            messages.append({
                "role": "assistant", "content": None,
                "tool_calls": [{"id": f"tc{i}", "function": {"name": "shell"}}],
            })
            messages.append({"role": "tool", "tool_call_id": f"tc{i}", "content": "r" * 200})
        cutoff = find_cutoff(messages, head_len=1)
        self.assertIsNotNone(cutoff)
        self.assertEqual(messages[cutoff]["role"], "user")


class UserOrdinalMappingTests(unittest.TestCase):
    def test_ordinal_roundtrip_across_prepared_list(self) -> None:
        """prepare 补占位 tool 消息后 index 平移,但「第 k 条 user」两边一致。"""
        messages = [
            {"role": "system", "content": "s"},
            {"role": "user", "content": "q1"},
            # 悬空 tool_call:prepare 的 repair 会在其后补一条占位 tool 消息
            {"role": "assistant", "content": None,
             "tool_calls": [{"id": "tc1", "function": {"name": "shell"}}]},
            {"role": "user", "content": "q2"},
            {"role": "assistant", "content": "a2"},
        ]
        from core.context import prepare_messages_with_stats
        prepared, stats = prepare_messages_with_stats(messages)
        self.assertEqual(stats["repaired_tool_calls"], 1)
        k = _user_ordinal(messages, 3)          # q2 是第 2 条 user
        self.assertEqual(k, 2)
        p_idx = _kth_user_index(prepared, k)
        self.assertEqual(prepared[p_idx]["content"], "q2")
        self.assertEqual(p_idx, 4)  # 占位 tool 插入使 index +1

    def test_kth_user_missing_returns_none(self) -> None:
        self.assertIsNone(_kth_user_index([{"role": "assistant", "content": "a"}], 1))


class ApplyFoldTests(unittest.TestCase):
    def test_apply_fold_rebuilds_head_and_advances_base(self) -> None:
        sess = Session(task_id=uuid4(), system_prompt="SYS")
        window = _mk_window(10)
        sess.messages.extend(window)          # 模拟 load:base=0,head=1(system)
        cutoff = 11                            # 内存 index,某条 user 消息
        self.assertEqual(sess.messages[cutoff]["role"], "user")

        sess.apply_fold(cutoff, "SUMMARY-TEXT")

        # 窗口 = [system] + [摘要注入] + 尾段
        self.assertEqual(sess.messages[0]["role"], "system")
        self.assertTrue(sess.messages[1]["content"].startswith(SUMMARY_NOTE_HEADER))
        self.assertIn("SUMMARY-TEXT", sess.messages[1]["content"])
        self.assertEqual(sess.messages[2]["content"], window[cutoff - 1]["content"])
        # base 推进 = 折叠条数(cutoff - 旧 head 1);head 变 2(system + 摘要)
        self.assertEqual(sess.context_base, cutoff - 1)
        self.assertEqual(sess.context_head_len, 2)

    def test_second_fold_replaces_summary_note(self) -> None:
        sess = Session(task_id=uuid4(), system_prompt="SYS")
        sess.messages.extend(_mk_window(10))
        sess.apply_fold(11, "FIRST")
        base_after_first = sess.context_base
        sess.messages.extend(_mk_window(10))
        # head=2,再折 10 条(切到 index 12 的 user 消息)
        self.assertEqual(sess.messages[12]["role"], "user")
        sess.apply_fold(12, "SECOND")
        notes = [m for m in sess.messages if str(m.get("content", "")).startswith(SUMMARY_NOTE_HEADER)]
        self.assertEqual(len(notes), 1)        # 旧摘要被折进新摘要,不叠加
        self.assertIn("SECOND", notes[0]["content"])
        self.assertEqual(sess.context_base, base_after_first + 10)
        self.assertEqual(sess.context_head_len, 2)

    def test_summary_note_not_counted_as_user_msg(self) -> None:
        sess = Session(task_id=uuid4(), system_prompt="SYS")
        sess.messages.extend(_mk_window(10))
        before = sess.n_user_msgs()
        sess.apply_fold(11, "S")
        # 折掉 5 轮 user(index 1..10 里 5 条 user),注入的摘要不计入
        self.assertEqual(sess.n_user_msgs(), before - 5)


class _FakeLLM:
    def __init__(self, reply: str) -> None:
        self.reply = reply
        self.last_messages = None

    def chat(self, messages, **kwargs):
        self.last_messages = messages
        return SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content=self.reply))],
            usage=SimpleNamespace(prompt_tokens=100, completion_tokens=20),
        )


_FAKE_CAPS = SimpleNamespace(
    reliable_context=100,  # 极小上限 → 体量必然超阈值,不用造几十万字符
    family="deepseek_v4", variant="flash",
    input_cny_per_mtoken=1.0, output_cny_per_mtoken=2.0,
)


class MaybeFoldTests(unittest.TestCase):
    def _session(self) -> Session:
        sess = Session(task_id=uuid4(), system_prompt="SYS")
        sess.messages.extend(_mk_window(30))
        return sess

    def test_fold_happy_path(self) -> None:
        sess = self._session()
        llm = _FakeLLM("## 目标\n...")
        events = []
        with patch.object(cf, "persist_fold") as persist, \
             patch.object(cf, "record_chat_usage") as usage:
            stats = cf.maybe_fold(sess, llm, _FAKE_CAPS, user_id=uuid4(), emit=events.append)
        self.assertIsNotNone(stats)
        self.assertEqual(stats["phase"], "done")
        # 摘要指令是折叠调用的最后一条消息(前缀复用构造)
        self.assertEqual(llm.last_messages[-1]["content"], cf.FOLD_INSTRUCTION)
        # DB 持久化先行,参数与内存变更一致
        persist.assert_called_once_with(sess.task_id, stats["new_base"], "## 目标\n...")
        usage.assert_called_once()
        # 独立 kind:折叠次数 = count(usage_events kind='context_fold'),task 详情接口依赖
        self.assertEqual(usage.call_args.kwargs.get("kind"), "context_fold")
        self.assertEqual(sess.context_base, stats["new_base"])
        phases = [e.get("phase") for e in events if e.get("type") == "context_fold"]
        self.assertEqual(phases, ["start", "done"])
        self.assertLess(stats["after_chars"], stats["before_chars"])

    def test_below_threshold_no_fold(self) -> None:
        sess = self._session()
        caps = SimpleNamespace(**{**_FAKE_CAPS.__dict__, "reliable_context": 10_000_000})
        llm = _FakeLLM("unused")
        with patch.object(cf, "persist_fold") as persist:
            self.assertIsNone(cf.maybe_fold(sess, llm, caps, user_id=uuid4(), emit=lambda e: None))
        persist.assert_not_called()
        self.assertIsNone(llm.last_messages)

    def test_empty_summary_skips_fold(self) -> None:
        sess = self._session()
        before = list(sess.messages)
        events = []
        with patch.object(cf, "persist_fold") as persist:
            result = cf.maybe_fold(
                sess, _FakeLLM("   "), _FAKE_CAPS, user_id=uuid4(), emit=events.append
            )
        self.assertIsNone(result)
        persist.assert_not_called()
        self.assertEqual(sess.messages, before)   # 内存零污染
        self.assertTrue(any(e.get("type") == "warn" for e in events))


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