"""
测试会话数据完整性 —— 单轨结构化格式全面验证

覆盖场景：
1. _save_session() 保存全部结构化字段（tool_calls / tool_call_id / reasoning_content）
2. 多轮任务间 tool 消息在 current_messages 中累积不丢失
3. get_messages() 在会话重载后返回完整数据
4. 旧版 structured_messages 双轨格式加载兼容
"""
import sys
import tempfile
import asyncio
import json
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import pytest
from aacode.utils.session_manager import SessionManager, SessionMessage


@pytest.fixture
def sm():
    with tempfile.TemporaryDirectory() as td:
        yield SessionManager(Path(td))


def _read_session_file(sm):
    """直接读取会话 JSON 文件，绕过 SessionManager 缓存"""
    f = sm.sessions_dir / f"{sm.current_session_id}.json"
    with open(f, "r") as fh:
        return json.load(fh)


class TestSaveStructuredFields:
    """_save_session 保存结构化字段到 JSON"""

    @pytest.mark.asyncio
    async def test_tool_calls_in_json(self, sm):
        sid = await sm.create_session("test")
        sm.current_messages.append(SessionMessage(
            role="assistant", content="checking",
            timestamp="2026-01-01T00:00:00",
            tool_calls=[{"name": "run_shell", "arguments": '{"cmd":"ls"}'}],
            reasoning_content="need to check",
        ))
        await sm._save_session()
        raw = _read_session_file(sm)
        msgs = raw["messages"]
        assistant_msgs = [m for m in msgs if m["role"] == "assistant"]
        assert len(assistant_msgs) >= 1
        assert assistant_msgs[-1]["tool_calls"] == [{"name": "run_shell", "arguments": '{"cmd":"ls"}'}]
        assert assistant_msgs[-1]["reasoning_content"] == "need to check"
        assert "structured_messages" not in raw

    @pytest.mark.asyncio
    async def test_tool_call_id_in_json(self, sm):
        await sm.create_session("test")
        sm.current_messages.append(SessionMessage(
            role="tool", content="result",
            timestamp="2026-01-01T00:00:00",
            tool_call_id="c1",
        ))
        await sm._save_session()
        raw = _read_session_file(sm)
        tool_msgs = [m for m in raw["messages"] if m["role"] == "tool"]
        assert len(tool_msgs) >= 1
        assert tool_msgs[-1]["tool_call_id"] == "c1"

    @pytest.mark.asyncio
    async def test_timestamp_and_tokens_present(self, sm):
        await sm.create_session("test")
        await sm.add_message("assistant", "hello")
        raw = _read_session_file(sm)
        last = raw["messages"][-1]
        assert "timestamp" in last
        assert "tokens" in last
        assert isinstance(last["timestamp"], str)
        assert isinstance(last["tokens"], int)


class TestToolMessagesAccumulateAcrossTasks:
    """模拟多轮任务，tool 消息累积不丢失"""

    @pytest.mark.asyncio
    async def test_tool_messages_accumulate(self, sm):
        await sm.create_session("test")

        # 任务 1
        sm.current_messages.append(SessionMessage(
            role="assistant", content="ok",
            timestamp="2026-01-01T00:00:00",
            tool_calls=[{"name": "run_shell"}],
        ))
        sm.current_messages.append(SessionMessage(
            role="tool", content="result_A",
            timestamp="2026-01-01T00:00:01",
            tool_call_id="c1",
        ))

        msgs_before = await sm.get_messages(include_system=False)
        tool_count_before = len([m for m in msgs_before if m["role"] == "tool"])
        assert tool_count_before == 1

        # 任务 2
        sm.current_messages.append(SessionMessage(
            role="assistant", content="ok",
            timestamp="2026-01-01T00:00:02",
            tool_calls=[{"name": "search"}],
        ))
        sm.current_messages.append(SessionMessage(
            role="tool", content="result_B",
            timestamp="2026-01-01T00:00:03",
            tool_call_id="c2",
        ))

        msgs_after = await sm.get_messages(include_system=False)
        tool_count_after = len([m for m in msgs_after if m["role"] == "tool"])
        assert tool_count_after == 2


class TestReloadSessionConsistency:
    """会话重载后数据一致性"""

    @pytest.mark.asyncio
    async def test_reload_gets_tool_messages(self, sm):
        sid = await sm.create_session("test")
        sm.current_messages.append(SessionMessage(
            role="user", content="q",
            timestamp="2026-01-01T00:00:00",
        ))
        sm.current_messages.append(SessionMessage(
            role="assistant", content="a",
            timestamp="2026-01-01T00:00:01",
            tool_calls=[],
        ))
        sm.current_messages.append(SessionMessage(
            role="tool", content="r",
            timestamp="2026-01-01T00:00:02",
            tool_call_id="c1",
        ))
        await sm._save_session()

        sm2 = SessionManager(sm.project_path)
        ok = await sm2.switch_session(sid)
        assert ok

        msgs = await sm2.get_messages(include_system=False)
        roles = [m["role"] for m in msgs]
        assert "user" in roles
        assert "assistant" in roles
        assert "tool" in roles

        tool_msgs = [m for m in msgs if m["role"] == "tool"]
        assert len(tool_msgs) == 1
        assert tool_msgs[0]["tool_call_id"] == "c1"

    @pytest.mark.asyncio
    async def test_add_message_after_reload(self, sm):
        sid = await sm.create_session("test")
        sm.current_messages.append(SessionMessage(
            role="user", content="original",
            timestamp="2026-01-01T00:00:00",
        ))
        sm.current_messages.append(SessionMessage(
            role="assistant", content="response",
            timestamp="2026-01-01T00:00:01",
        ))
        await sm._save_session()

        sm2 = SessionManager(sm.project_path)
        await sm2.switch_session(sid)

        await sm2.add_message("user", "new question")
        msgs = await sm2.get_messages(include_system=False)
        contents = [m["content"] for m in msgs if m["role"] == "user"]
        assert any("original" in c for c in contents)
        assert any("new question" in c for c in contents)

    @pytest.mark.asyncio
    async def test_old_dual_track_format_compat(self, sm):
        """旧版 structured_messages 双轨文件加载后，数据完整"""
        sid = await sm.create_session("test")
        # 模拟旧版双轨文件：同时有 messages 和 structured_messages
        await sm._save_session()
        raw = _read_session_file(sm)
        raw["structured_messages"] = [
            {"role": "user", "content": "q"},
            {"role": "assistant", "content": "a", "tool_calls": [{"name": "run_shell"}]},
            {"role": "tool", "content": "r", "tool_call_id": "c1"},
        ]
        with open(sm.sessions_dir / f"{sid}.json", "w") as f:
            json.dump(raw, f)

        sm2 = SessionManager(sm.project_path)
        ok = await sm2.switch_session(sid)
        assert ok

        msgs = await sm2.get_messages(include_system=False)
        roles = [m["role"] for m in msgs]
        assert "user" in roles
        assert "assistant" in roles
        assert "tool" in roles
        tool_msgs = [m for m in msgs if m["role"] == "tool"]
        assert len(tool_msgs) == 1
        assert tool_msgs[0]["tool_call_id"] == "c1"


class TestGetSessionConsistency:
    """get_session API 返回完整数据"""

    @pytest.mark.asyncio
    async def test_get_session_returns_all_roles(self, sm):
        sid = await sm.create_session("test")
        sm.current_messages.append(SessionMessage(
            role="user", content="hello",
            timestamp="2026-01-01T00:00:00",
        ))
        sm.current_messages.append(SessionMessage(
            role="assistant", content="hi",
            timestamp="2026-01-01T00:00:01",
            tool_calls=[{"name": "run_shell"}],
        ))
        sm.current_messages.append(SessionMessage(
            role="tool", content="output",
            timestamp="2026-01-01T00:00:02",
            tool_call_id="c1",
        ))
        await sm._save_session()

        session_data = await sm.get_session(sid)
        assert "messages" in session_data
        msgs = session_data["messages"]
        roles = [m["role"] for m in msgs]
        assert "user" in roles
        assert "assistant" in roles
        assert "tool" in roles

    @pytest.mark.asyncio
    async def test_get_session_after_reload(self, sm):
        sid = await sm.create_session("test")
        await sm.add_message("user", "first msg")
        await sm.add_message("assistant", "first resp")
        await sm._save_session()

        sm2 = SessionManager(sm.project_path)
        await sm2.switch_session(sid)
        session_data = await sm2.get_session(sid)
        msgs = session_data["messages"]
        assert len(msgs) >= 2
        assert any("first msg" in m.get("content", "") for m in msgs)
