# -*- coding: utf-8 -*-
"""
会话完整性集成测试：覆盖保存/恢复/上下文/显示 4 个关键链路。

发现的 Bug 集合：
- A: _print_resumed_messages 跳过 content="" 的 assistant（含 tool_calls）
     → 用户视觉上"看不到历史 assistant 消息"
- B: _sanitize_messages 在 messages[0] 非 system 时丢失 system prompt
     → 恢复后 LLM 调用缺 system，无系统指令
- C: _sanitize_messages 在 messages 为空时返回空，session 缺 system
- D: 上下文压缩替换 session.messages 时不应丢失 system 首条
"""
import json
from pathlib import Path

import pytest

from abyss.session import Session


SYSTEM = "你是 abyss 助手"


def _tc(idx: str, name: str = "todo_write"):
    return {
        "id": idx, "type": "function",
        "function": {"name": name, "arguments": "{}"},
    }


def _write(tmp_path: Path, messages: list, name: str = "latest") -> Path:
    p = tmp_path / f"{name}.json"
    p.write_text(json.dumps({
        "messages": messages,
        "tool_call_counter": 0,
    }, ensure_ascii=False), encoding="utf-8")
    return p


# ---------- Bug B: system 丢失 ----------

class TestSystemPreservation:
    """验证 from_file 在任何情况下都保证 messages[0] 是 system。"""

    def test_system_lost_due_to_truncation(self, tmp_path):
        """场景：messages[0] 就是不合法 assistant，截断后应插入 system 兜底。"""
        msgs = [
            {"role": "assistant", "content": "", "tool_calls": [_tc("c1")]},
            # 没 tool 响应，无 system
        ]
        p = _write(tmp_path, msgs)

        s = Session.from_file(p, SYSTEM)

        assert len(s.messages) >= 1
        assert s.messages[0]["role"] == "system"
        assert s.messages[0]["content"] == SYSTEM

    def test_empty_messages_gets_system(self, tmp_path):
        """空 messages 列表应至少返回 system。"""
        p = _write(tmp_path, [])

        s = Session.from_file(p, SYSTEM)

        assert len(s.messages) == 1
        assert s.messages[0]["role"] == "system"

    def test_messages_without_system_at_start(self, tmp_path):
        """messages[0] 是 user 时，恢复后应在首位补 system。"""
        msgs = [
            {"role": "user", "content": "hi"},
            {"role": "assistant", "content": "hello"},
        ]
        p = _write(tmp_path, msgs)

        s = Session.from_file(p, SYSTEM)

        assert s.messages[0]["role"] == "system"
        assert s.messages[0]["content"] == SYSTEM
        # 后续消息保留
        assert s.messages[1]["role"] == "user"
        assert s.messages[2]["role"] == "assistant"

    def test_normal_session_preserves_system(self, tmp_path):
        """正常会话：messages[0] 已是 system，不应被替换。"""
        original_system = "原始 system prompt"
        msgs = [
            {"role": "system", "content": original_system},
            {"role": "user", "content": "hi"},
        ]
        p = _write(tmp_path, msgs)

        s = Session.from_file(p, SYSTEM)

        # 应保留原 system，不覆盖
        assert s.messages[0]["role"] == "system"
        assert s.messages[0]["content"] == original_system


# ---------- Bug A: 显示 bug ----------

class TestResumedMessagesDisplay:
    """验证 _print_resumed_messages 显示所有 assistant 消息（含 tool_calls）。"""

    def test_display_assistant_with_only_tool_calls(self, tmp_path, capsys):
        """assistant 只有 tool_calls、content="" 时也应显示占位。"""
        from abyss.main import _print_resumed_messages

        session = Session(SYSTEM)
        session.messages = [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": "帮我做任务"},
            {"role": "assistant", "content": "", "tool_calls": [_tc("c1", "todo_write")]},
            {"role": "tool", "tool_call_id": "c1", "content": "done"},
        ]

        _print_resumed_messages(session)

        out = capsys.readouterr().out
        # 应该显示 user 消息
        assert "帮我做任务" in out
        # 应该显示 assistant（即使是占位）
        assert "todo_write" in out or "工具" in out or "[调用" in out

    def test_display_normal_assistant(self, tmp_path, capsys):
        """正常 assistant 应显示 content。"""
        from abyss.main import _print_resumed_messages

        session = Session(SYSTEM)
        session.messages = [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": "你好"},
            {"role": "assistant", "content": "你好，有什么可以帮你？"},
        ]

        _print_resumed_messages(session)

        out = capsys.readouterr().out
        assert "你好，有什么可以帮你？" in out


# ---------- 端到端：保存 → 恢复 → 上下文完整 ----------

class TestSaveRestoreRoundTrip:
    """端到端：保存的会话能完整恢复，且可作为 LLM 上下文使用。"""

    def test_round_trip_preserves_all_messages(self, tmp_path):
        """保存的会话恢复后应包含所有原始消息。"""
        original = Session(SYSTEM)
        original.add_user("第一问")
        original.add_assistant(content="第一答")
        original.add_user("第二问")
        original.add_assistant(
            content="",
            tool_calls=[_tc("c1", "file_read")]
        )
        original.add_tool_result("c1", "文件内容")

        path = tmp_path / "round.json"
        original.save(path)

        restored = Session.from_file(path, SYSTEM)

        assert len(restored.messages) == len(original.messages)
        assert [m["role"] for m in restored.messages] == [
            "system", "user", "assistant", "user", "assistant", "tool"
        ]
        # tool_call_id 和 tool_calls 都应保留
        assert restored.messages[4]["tool_calls"][0]["id"] == "c1"
        assert restored.messages[5]["tool_call_id"] == "c1"

    def test_restored_session_can_provide_api_messages(self, tmp_path):
        """恢复的会话应能产出合法的 API messages。"""
        original = Session(SYSTEM)
        original.add_user("你好")
        original.add_assistant(content="你好啊")

        path = tmp_path / "api.json"
        original.save(path)

        restored = Session.from_file(path, SYSTEM)
        api_messages = restored.get_messages_for_api()

        # API messages 必须以 system 开头
        assert api_messages[0]["role"] == "system"
        # 不能为空
        assert len(api_messages) >= 3
        # 不能有 reasoning_content
        for m in api_messages:
            assert "reasoning_content" not in m


# ---------- Bug C: 上下文压缩不应丢失 system ----------

class TestContextCompressionPreservesSystem:
    """验证上下文压缩替换 session.messages 后 system 首条不丢失。"""

    def test_compress_keeps_system_first(self):
        """模拟压缩引擎返回不含 system 的列表，应被恢复。"""
        session = Session(SYSTEM)
        session.add_user("hi")
        session.add_assistant(content="hello")

        # 模拟压缩引擎 bug：返回不含 system 的列表
        compressed = [
            {"role": "user", "content": "compressed hi"},
            {"role": "assistant", "content": "compressed hello"},
        ]

        # 这里直接替换，模拟 main.py:1205 的 session.messages = compressed
        # 但应有守卫保证 system 不丢
        # 当前代码无守卫，所以这个测试会失败
        from abyss.session import Session as S
        # 期望有新方法 _ensure_system_preserved
        if hasattr(S, "_ensure_system_preserved"):
            session.messages = S._ensure_system_preserved(compressed, SYSTEM)
        else:
            session.messages = compressed  # 旧逻辑

        # 守卫应保证 system 在首位
        assert session.messages[0]["role"] == "system"


# ---------- 孤儿 assistant(tool_calls) 消息防御 ----------
# 场景：agent_loop.process_tool_calls 在 add_assistant(tool_calls=...) 之后
# 工具执行被 TaskInterrupted 中断，若不回滚会留下孤儿消息，导致下一次
# API 调用报 400 "insufficient tool messages following tool_calls"。


def test_get_messages_for_api_truncates_orphan_assistant(tmp_path):
    """孤儿 assistant(tool_calls) 后无 tool 结果时，get_messages_for_api 应截断。"""
    session = Session(SYSTEM)
    session.add_user("do something")
    # 模拟中断：assistant 写入了 tool_calls，但无 tool 结果
    session.add_assistant(tool_calls=[_tc("call_1")])

    result = session.get_messages_for_api()

    # 孤儿 assistant 应被截断
    roles = [m["role"] for m in result]
    assert "assistant" not in roles or not result[-1].get("tool_calls"), \
        "孤儿 assistant(tool_calls) 消息应被 get_messages_for_api 截断"
    # session.messages 也应被同步清理
    assert not any(
        m.get("role") == "assistant" and m.get("tool_calls")
        and not _has_following_tool_result(session.messages, i)
        for i, m in enumerate(session.messages)
    ), "session.messages 不应残留孤儿 assistant(tool_calls)"


def _has_following_tool_result(messages, idx):
    """检查 messages[idx] 后是否紧跟覆盖所有 tool_call_id 的 tool 消息。"""
    msg = messages[idx]
    tool_calls = msg.get("tool_calls") or []
    expected = {tc.get("id") for tc in tool_calls if tc.get("id")}
    got = set()
    j = idx + 1
    while j < len(messages) and messages[j].get("role") == "tool":
        tc_id = messages[j].get("tool_call_id")
        if tc_id:
            got.add(tc_id)
        j += 1
    return expected.issubset(got)


def test_rollback_last_assistant_removes_orphan():
    """rollback_last_assistant 应移除末尾的孤儿 assistant 消息。"""
    session = Session(SYSTEM)
    session.add_user("hi")
    session.add_assistant(tool_calls=[_tc("call_1")])

    assert session.messages[-1]["role"] == "assistant"
    session.rollback_last_assistant()
    assert session.messages[-1]["role"] == "user"


def test_rollback_last_assistant_does_not_remove_tool_message():
    """末尾是 tool 消息时不应回滚（已有 tool 结果，不是孤儿）。"""
    session = Session(SYSTEM)
    session.add_user("hi")
    session.add_assistant(tool_calls=[_tc("call_1")])
    session.add_tool_result("call_1", '{"success": true}')

    session.rollback_last_assistant()
    # tool 消息应保留
    assert session.messages[-1]["role"] == "tool"


def test_rollback_last_assistant_empty_session_safe():
    """空 session 或只有 system 时调用应安全无操作。"""
    session = Session(SYSTEM)
    session.rollback_last_assistant()
    assert len(session.messages) == 1  # 仅 system


def test_get_messages_for_api_keeps_valid_tool_sequence():
    """完整的 assistant(tool_calls) + tool 结果序列不应被截断。"""
    session = Session(SYSTEM)
    session.add_user("do something")
    session.add_assistant(tool_calls=[_tc("call_1"), _tc("call_2")])
    session.add_tool_result("call_1", '{"success": true}')
    session.add_tool_result("call_2", '{"success": true}')

    result = session.get_messages_for_api()
    # 完整序列应保留
    roles = [m["role"] for m in result]
    assert roles == ["system", "user", "assistant", "tool", "tool"]
