"""
近期改动专项测试（2026-05-27 ~ 2026-05-29）

覆盖：
1. _atomic_file_write 原子写入、错误清理
2. _merge_sessions_index 合并逻辑、并发不丢失
3. _save_iteration_messages 增量持久化
4. _finalize_session_save 最终保存
5. react_loop on_new_messages 回调完整性
6. _format_observation_for_display run_shell 截断
7. reasoning_content 与 content 分离存储
8. _save_session + _save_sessions_index merge 行为
"""
import sys
import json
import tempfile
from pathlib import Path
from datetime import datetime

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

import pytest

from aacode.utils.session_manager import (
    SessionManager, SessionMessage,
    _atomic_file_write, _merge_sessions_index,
)
from aacode.core.main_agent import MainAgent


# ============================================================================
# 辅助函数
# ============================================================================

def _build_agent(sm, project_path=None):
    """构造最小 MainAgent，注入 session_manager"""
    from aacode.utils.context_manager import ContextManager
    from aacode.utils.safety import SafetyGuard

    pp = Path(project_path or tempfile.mkdtemp())
    cm = ContextManager(pp)
    sg = SafetyGuard(pp)
    agent = MainAgent(
        project_path=pp,
        context_manager=cm,
        safety_guard=sg,
        model_config={"name": "mock-model", "base_url": "http://mock", "gateway": "openai"},
    )
    agent.session_manager = sm
    agent.system_prompt = "You are a coding assistant."
    agent.conversation_history = [{"role": "system", "content": agent.system_prompt}]
    return agent


async def _create_session_and_agent():
    tmp = Path(tempfile.mkdtemp())
    sm = SessionManager(tmp)
    await sm.create_session("test_task")
    agent = _build_agent(sm, tmp)
    return sm, agent


# ============================================================================
#  1. _atomic_file_write 测试
# ============================================================================

class TestAtomicFileWrite:
    """原子写入行为"""

    def test_basic_write(self, tmp_path):
        filepath = tmp_path / "test.json"
        _atomic_file_write(filepath, lambda f: f.write("hello"))
        assert filepath.read_text() == "hello"

    def test_overwrite_existing(self, tmp_path):
        filepath = tmp_path / "test.json"
        filepath.write_text("old")
        _atomic_file_write(filepath, lambda f: f.write("new"))
        assert filepath.read_text() == "new"

    def test_large_content(self, tmp_path):
        filepath = tmp_path / "large.json"
        large_data = "x" * 1_000_000
        _atomic_file_write(filepath, lambda f: f.write(large_data))
        assert len(filepath.read_text()) == 1_000_000

    def test_unicode_content(self, tmp_path):
        filepath = tmp_path / "unicode.json"
        unicode_data = "你好世界\nこんにちは\n🌟✨"
        _atomic_file_write(filepath, lambda f: f.write(unicode_data))
        assert filepath.read_text() == unicode_data

    def test_temp_file_cleaned_on_write_error(self, tmp_path):
        def failing_write(f):
            raise IOError("Disk full")

        filepath = tmp_path / "test.json"
        initial_tmp = len(list(tmp_path.glob("*.tmp")))
        with pytest.raises(IOError):
            _atomic_file_write(filepath, failing_write)
        final_tmp = len(list(tmp_path.glob("*.tmp")))
        assert final_tmp == initial_tmp

    def test_empty_write(self, tmp_path):
        filepath = tmp_path / "empty.json"
        _atomic_file_write(filepath, lambda f: f.write(""))
        assert filepath.read_text() == ""

    def test_json_dump_via_atomic_write(self, tmp_path):
        filepath = tmp_path / "data.json"
        data = {"key": "value", "list": [1, 2, 3]}
        _atomic_file_write(filepath, lambda f: json.dump(data, f))
        loaded = json.loads(filepath.read_text())
        assert loaded == data

    def test_newline_handling(self, tmp_path):
        filepath = tmp_path / "multiline.json"
        data = "line1\nline2\nline3\n"
        _atomic_file_write(filepath, lambda f: f.write(data))
        assert filepath.read_text() == data


# ============================================================================
#  2. _merge_sessions_index 测试
# ============================================================================

class TestMergeSessionsIndex:
    """合并会话索引"""

    def test_merge_new_data_into_empty_file(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        our_data = {"session_1": {"id": "1", "title": "test1"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert loaded["session_1"]["title"] == "test1"

    def test_merge_preserves_disk_entries(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        filepath.write_text(json.dumps({
            "session_1": {"id": "1", "title": "disk1"},
            "session_2": {"id": "2", "title": "disk2"},
        }))
        our_data = {"session_3": {"id": "3", "title": "new3"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert "session_1" in loaded
        assert "session_2" in loaded
        assert "session_3" in loaded
        assert loaded["session_1"]["title"] == "disk1"

    def test_our_data_overrides_disk_same_key(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        filepath.write_text(json.dumps({
            "session_1": {"id": "1", "title": "old_title", "status": "active"},
        }))
        our_data = {"session_1": {"id": "1", "title": "new_title"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert loaded["session_1"]["title"] == "new_title"
        assert "status" not in loaded["session_1"]

    def test_merge_corrupted_disk_file(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        filepath.write_text("{ corrupted json !!! }")
        our_data = {"session_1": {"id": "1"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert loaded["session_1"]["id"] == "1"

    def test_merge_missing_disk_file(self, tmp_path):
        filepath = tmp_path / "nonexistent.json"
        our_data = {"session_1": {"id": "1"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert loaded["session_1"]["id"] == "1"

    def test_merge_empty_disk_file(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        filepath.write_text("")
        our_data = {"session_1": {"id": "1"}}
        _merge_sessions_index(filepath, our_data)
        loaded = json.loads(filepath.read_text())
        assert loaded["session_1"]["id"] == "1"

    def test_merge_empty_our_data_preserves_disk(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        disk = {"session_1": {"id": "1", "title": "keep_me"}}
        filepath.write_text(json.dumps(disk))
        _merge_sessions_index(filepath, {})
        loaded = json.loads(filepath.read_text())
        assert loaded == disk

    def test_merge_large_index(self, tmp_path):
        filepath = tmp_path / "sessions_index.json"
        disk_data = {f"session_{i}": {"id": str(i), "title": f"disk_{i}"} for i in range(100)}
        filepath.write_text(json.dumps(disk_data))
        new_data = {f"session_{i+90}": {"id": str(i+90), "title": f"new_{i+90}"} for i in range(20)}
        _merge_sessions_index(filepath, new_data)
        loaded = json.loads(filepath.read_text())
        assert len(loaded) >= 100
        # 重叠的 key 被覆盖
        assert loaded["session_90"]["title"] == "new_90"


# ============================================================================
#  3. _save_iteration_messages 增量持久化
# ============================================================================

class TestSaveIterationMessages:
    """增量消息保存"""

    @pytest.mark.asyncio
    async def test_saves_assistant_messages(self):
        sm, agent = await _create_session_and_agent()
        msg_before = len(sm.current_messages)

        await agent._save_iteration_messages([
            {"role": "assistant", "content": "I'll run a command",
             "tool_calls": [{"id": "c1", "type": "function",
                             "function": {"name": "run_shell", "arguments": '{"command":"ls"}'}}],
             "reasoning_content": "need to list files"},
            {"role": "tool", "content": "file1.py\nfile2.py", "tool_call_id": "c1"},
        ])

        assert len(sm.current_messages) == msg_before + 2
        roles = [m.role for m in sm.current_messages[-2:]]
        assert roles == ["assistant", "tool"]

    @pytest.mark.asyncio
    async def test_skips_user_and_system_messages(self):
        sm, agent = await _create_session_and_agent()
        msg_before = len(sm.current_messages)

        await agent._save_iteration_messages([
            {"role": "system", "content": "extra system msg"},
            {"role": "user", "content": "extra user msg"},
            {"role": "assistant", "content": "only this should be saved"},
        ])

        assert len(sm.current_messages) == msg_before + 1
        assert sm.current_messages[-1].role == "assistant"

    @pytest.mark.asyncio
    async def test_handles_empty_new_msgs(self):
        sm, agent = await _create_session_and_agent()
        msg_before = len(sm.current_messages)
        await agent._save_iteration_messages([])
        assert len(sm.current_messages) == msg_before

    @pytest.mark.asyncio
    async def test_saves_full_reasoning_content(self):
        sm, agent = await _create_session_and_agent()
        reasoning = "I need to think about this very carefully..."
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "done", "reasoning_content": reasoning}
        ])
        assert sm.current_messages[-1].reasoning_content == reasoning

    @pytest.mark.asyncio
    async def test_saves_tool_call_id(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "tool", "content": "result", "tool_call_id": "call_abc123"}
        ])
        tool_msgs = [m for m in sm.current_messages if m.role == "tool"]
        assert tool_msgs[-1].tool_call_id == "call_abc123"

    @pytest.mark.asyncio
    async def test_multiple_iterations_accumulate(self):
        sm, agent = await _create_session_and_agent()
        for i in range(3):
            await agent._save_iteration_messages([
                {"role": "assistant", "content": f"step {i}"},
                {"role": "tool", "content": f"result {i}", "tool_call_id": f"c{i}"},
            ])
        assistant_count = sum(1 for m in sm.current_messages if m.role == "assistant")
        tool_count = sum(1 for m in sm.current_messages if m.role == "tool")
        assert assistant_count >= 3
        assert tool_count >= 3

    @pytest.mark.asyncio
    async def test_updates_session_index_after_save(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "hello"}
        ])
        summary = sm.sessions_index.get(sm.current_session_id)
        assert summary is not None
        assert summary.total_messages >= 2

    @pytest.mark.asyncio
    async def test_does_not_crash_on_missing_session(self):
        tmp = Path(tempfile.mkdtemp())
        sm = SessionManager(tmp)
        agent = _build_agent(sm, tmp)
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "orphan message"}
        ])

    @pytest.mark.asyncio
    async def test_session_file_persisted_to_disk(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "disk persisted"}
        ])
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        contents = [m["content"] for m in data["messages"]]
        assert "disk persisted" in contents


# ============================================================================
#  4. _finalize_session_save 测试
# ============================================================================

class TestFinalizeSessionSave:
    """最终会话保存"""

    @pytest.mark.asyncio
    async def test_saves_all_messages(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "step 1"},
            {"role": "tool", "content": "result 1", "tool_call_id": "c1"},
        ])
        msg_before = len(sm.current_messages)
        await agent._finalize_session_save()
        assert len(sm.current_messages) == msg_before

    @pytest.mark.asyncio
    async def test_updates_metadata(self):
        sm, agent = await _create_session_and_agent()
        sid = sm.current_session_id
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "hello"}
        ])
        await agent._finalize_session_save()
        summary = sm.sessions_index[sid]
        assert summary.total_messages == len(sm.current_messages)
        assert summary.last_activity is not None

    @pytest.mark.asyncio
    async def test_no_session_id_handled_gracefully(self):
        tmp = Path(tempfile.mkdtemp())
        sm = SessionManager(tmp)
        agent = _build_agent(sm, tmp)
        await agent._finalize_session_save()

    @pytest.mark.asyncio
    async def test_double_finalize_idempotent(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "hello"}
        ])
        await agent._finalize_session_save()
        msg_before = len(sm.current_messages)
        await agent._finalize_session_save()
        assert len(sm.current_messages) == msg_before


# ============================================================================
#  5. react_loop on_new_messages 回调
# ============================================================================

class TestReactLoopOnNewMessages:
    """React 循环增量回调"""

    def test_on_new_messages_param_accepted(self):
        import inspect
        from aacode.core.react_loop import AsyncReActLoop
        sig = inspect.signature(AsyncReActLoop.run)
        params = list(sig.parameters.keys())
        assert "on_new_messages" in params

    def test_run_method_accepts_kwargs(self):
        """run 方法能接受 on_new_messages 关键字参数"""
        import inspect
        from aacode.core.react_loop import AsyncReActLoop
        sig = inspect.signature(AsyncReActLoop.run)
        params = sig.parameters
        assert "on_new_messages" in params
        param = params["on_new_messages"]
        assert param.default is None


# ============================================================================
#  6. _execute_action_internal run_shell JSON 格式
# ============================================================================

class TestExecuteActionInternal:
    """run_shell 返回格式验证"""

    def test_run_shell_truncation_code_exists(self):
        """验证 _execute_action_internal 中有 run_shell 截断逻辑"""
        import inspect
        from aacode.core.react_loop import AsyncReActLoop

        src = inspect.getsource(AsyncReActLoop._execute_action_internal)
        assert "json.dumps(summary, ensure_ascii=False)" in src
        assert 'summary.pop("command", "")' in src
        assert 'summary.pop("working_directory", None)' in src
        assert 'cmd.split(\'\\n\')[0][:200]' in src or "cmd.split(" in src

    def test_run_shell_no_double_command(self):
        """run_shell 结果不重复回传完整 command"""
        import inspect
        from aacode.core.react_loop import AsyncReActLoop

        src = inspect.getsource(AsyncReActLoop._execute_action_internal)
        assert "避免 command 重复" in src or "无需在 tool 消息中再次回传" in src


# ============================================================================
#  7. reasoning_content 与 content 分离
# ============================================================================

class TestReasoningContentSeparation:
    """推理内容与正文分离存储"""

    @pytest.mark.asyncio
    async def test_reasoning_saved_separately(self):
        sm, agent = await _create_session_and_agent()

        sm.current_messages.append(SessionMessage(
            role="assistant",
            content="The answer is 42",
            reasoning_content="I need to calculate the meaning of life",
            timestamp=datetime.now().isoformat(timespec='seconds'),
        ))
        await sm._save_session()

        sm2 = SessionManager(sm.project_path)
        await sm2.switch_session(sm.current_session_id)
        msgs = await sm2.get_messages(include_system=False)

        assistant = [m for m in msgs if m["role"] == "assistant"][-1]
        assert assistant["reasoning_content"] is not None
        assert "meaning of life" in assistant["reasoning_content"]
        assert "The answer is 42" in assistant["content"]
        assert "meaning of life" not in assistant["content"]


# ============================================================================
#  9. 会话数据完整性
# ============================================================================

class TestSessionDataIntegrity:
    """会话数据完整性"""

    @pytest.mark.asyncio
    async def test_total_messages_correct_after_incremental_save(self):
        sm, agent = await _create_session_and_agent()
        for i in range(5):
            await agent._save_iteration_messages([
                {"role": "assistant", "content": f"msg {i}"}
            ])
        summary = sm.sessions_index[sm.current_session_id]
        assert summary.total_messages == len(sm.current_messages)
        # user + 5 assistants = 6 (no placeholder system)
        assert summary.total_messages == 6

    @pytest.mark.asyncio
    async def test_session_file_json_valid(self):
        sm, agent = await _create_session_and_agent()
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        assert "session_id" in data
        assert "messages" in data
        assert isinstance(data["messages"], list)

    @pytest.mark.asyncio
    async def test_no_metadata_field_in_saved_messages(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "test"}
        ])
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        for msg in data["messages"]:
            assert "metadata" not in msg, f"metadata should not exist in {msg.get('role')}"

    @pytest.mark.asyncio
    async def test_no_structured_messages_key(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "hello"},
            {"role": "tool", "content": "ok", "tool_call_id": "c1"},
        ])
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        assert "structured_messages" not in data

    @pytest.mark.asyncio
    async def test_all_timestamps_are_iso_strings(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "test"}
        ])
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        for msg in data["messages"]:
            ts = msg.get("timestamp", "")
            assert isinstance(ts, str), f"timestamp should be str, got {type(ts)}: {ts}"

    @pytest.mark.asyncio
    async def test_tokens_field_is_int(self):
        sm, agent = await _create_session_and_agent()
        await agent._save_iteration_messages([
            {"role": "assistant", "content": "test"}
        ])
        session_file = sm.sessions_dir / f"{sm.current_session_id}.json"
        data = json.loads(session_file.read_text())
        for msg in data["messages"]:
            assert isinstance(msg.get("tokens"), int), f"tokens should be int in {msg.get('role')}"

    @pytest.mark.asyncio
    async def test_no_duplicate_system_messages(self):
        """增量保存不应产生重复的 system 消息"""
        sm, agent = await _create_session_and_agent()
        system_count_before = sum(1 for m in sm.current_messages if m.role == "system")

        for _ in range(3):
            await agent._save_iteration_messages([
                {"role": "assistant", "content": "msg"}
            ])

        system_count_after = sum(1 for m in sm.current_messages if m.role == "system")
        assert system_count_after == system_count_before


# ============================================================================
#  10. _save_session + index merge 行为
# ============================================================================

class TestSaveSessionMerge:
    """_save_session + _save_sessions_index 合并行为"""

    @pytest.mark.asyncio
    async def test_index_reflects_all_created_sessions(self):
        tmp = Path(tempfile.mkdtemp())
        sm1 = SessionManager(tmp)
        await sm1.create_session("task_1")

        sm2 = SessionManager(tmp)
        sid2 = await sm2.create_session("task_2")

        # sm1 再保存索引，不应丢失 sm2 的会话
        sm1._save_sessions_index()

        sm3 = SessionManager(tmp)
        sessions = await sm3.list_sessions()
        sids = [s["session_id"] for s in sessions]
        assert sm1.current_session_id in sids
        assert sid2 in sids

    @pytest.mark.asyncio
    async def test_concurrent_process_writes_not_lost(self):
        tmp = Path(tempfile.mkdtemp())

        # 模拟进程 A 和进程 B 都创建会话
        sm_a = SessionManager(tmp)
        sm_b = SessionManager(tmp)

        sid_a = await sm_a.create_session("task_a")
        sid_b = await sm_b.create_session("task_b")

        sm_a._save_sessions_index()
        sm_b._save_sessions_index()

        sm_c = SessionManager(tmp)
        sessions = await sm_c.list_sessions()
        sids = [s["session_id"] for s in sessions]
        assert sid_a in sids, "Session from process A should not be lost"
        assert sid_b in sids, "Session from process B should not be lost"
