"""
测试异常恢复 —— react_loop.run() 抛异常时数据不丢失

场景：
1. react_loop 正常完成 → 所有数据完整保存
2. react_loop 抛出非 CancelledError → 已执行 steps 仍被保存
3. react_loop 抛出 CancelledError → 不保存（向上传播）
4. react_loop 第一轮就异常 → user task 消息仍保存
"""
import sys
import asyncio
import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

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

import pytest
from aacode.utils.session_manager import SessionManager


class FakeContextManager:
    async def get_compact_context(self):
        return "fake context"
    async def update(self, *a, **kw):
        pass
    async def save_history(self, *a, **kw):
        return "/fake/path"


class FakeSafetyGuard:
    pass


class FakeReActLoop:
    """模拟 react_loop，可控制是否抛异常"""

    def __init__(self, fail_mode=None):
        self.steps = []
        self.last_messages = None
        self.fail_mode = fail_mode

    async def run(self, initial_prompt="", task_description="", todo_manager=None, history_messages=None, on_new_messages=None, session_id=None):
        step = MagicMock()
        step.raw_response = "💭 Thinking: let me check\n\nThought: I will run a command"
        step.thought = "I will run a command"
        step.actions = [
            MagicMock(
                action="run_shell",
                action_input={"cmd": "echo hello"},
                observation="hello\n",
            )
        ]
        self.steps = [step]
        new_msgs = [
            {"role": "assistant", "content": "thinking", "tool_calls": [{"name": "run_shell"}]},
            {"role": "tool", "content": "hello", "tool_call_id": "c1"},
            {"role": "assistant", "content": "Done"},
        ]
        self.last_messages = [
            {"role": "system", "content": initial_prompt},
            {"role": "user", "content": task_description},
        ]
        if history_messages:
            self.last_messages.extend(history_messages)
        self.last_messages.extend(new_msgs)

        if on_new_messages:
            await on_new_messages(new_msgs)

        if self.fail_mode == "exception":
            raise RuntimeError("Simulated task failure")
        elif self.fail_mode == "cancelled":
            raise asyncio.CancelledError("Cancelled by user")

        return {
            "status": "completed",
            "final_thought": "Done",
            "iterations": 1,
            "steps": self.steps,
            "total_time": 0.5,
        }


def _build_agent(sm, project_path, react_loop):
    """构造一个 MainAgent，注入 mock 依赖 + 给定的 react_loop"""
    from aacode.core.main_agent import MainAgent

    agent = MainAgent(
        project_path=Path(project_path),
        context_manager=FakeContextManager(),
        safety_guard=FakeSafetyGuard(),
        model_config={"name": "mock-model", "base_url": "http://mock"},
    )
    # 替换 session_manager 和 react_loop
    agent.session_manager = sm
    agent.react_loop = react_loop
    agent.system_prompt = "You are a coding assistant."
    agent.conversation_history = [{"role": "system", "content": agent.system_prompt}]
    return agent


@pytest.mark.asyncio
async def test_normal_execution_saves_all():
    """正常执行：所有消息完整保存"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        sid = await sm.create_session("prev task")
        await sm.add_message("user", "prev msg")
        await sm.add_message("assistant", "prev resp")

        agent = _build_agent(sm, td, FakeReActLoop(fail_mode=None))
        result = await agent.execute(
            task="do something",
            init_instructions="test instructions",
            max_iterations=10,
        )

        assert result["status"] == "completed"
        assert sm.current_session_id == sid

        all_msgs = sm.current_messages
        contents = [m.content for m in all_msgs]
        assert "prev msg" in contents
        assert "prev resp" in contents
        assert "do something" in contents
        assert any("✅ Task completed" in c for c in contents)

        sm2 = SessionManager(project_path)
        await sm2.switch_session(sid)
        msgs2 = await sm2.get_messages(include_system=False)
        c2 = [m["content"] for m in msgs2]
        assert "do something" in c2

        raw = json.loads(Path(sm.sessions_dir / f"{sid}.json").read_text())
        raw_msgs = [m["content"] for m in raw["messages"]]
        assert "prev msg" in raw_msgs
        assert "do something" in raw_msgs


@pytest.mark.asyncio
async def test_exception_saves_partial_data():
    """异常：已执行的 steps 仍保存，不加 ✅ 标记"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        sid = await sm.create_session("prev task")
        await sm.add_message("user", "prev msg")
        await sm.add_message("assistant", "prev resp")

        agent = _build_agent(sm, td, FakeReActLoop(fail_mode="exception"))
        result = await agent.execute(
            task="do something",
            init_instructions="test instructions",
            max_iterations=10,
        )

        assert result["status"] == "error"
        assert "Simulated task failure" in result["error"]
        assert result["iterations"] == 1
        assert result["session_id"] == sid

        all_msgs = sm.current_messages
        contents = [m.content for m in all_msgs]
        assert "prev msg" in contents
        assert "prev resp" in contents
        assert "do something" in contents
        assert not any("✅ Task completed" in c for c in contents)

        # 结构化消息保存在 current_messages 中（含来自 last_messages 的内容）
        contents = [m.content for m in sm.current_messages]
        assert any("Done" in c for c in contents)


@pytest.mark.asyncio
async def test_exception_first_step_only_saves_user_msg():
    """第一轮就异常：仅保存 user task 消息"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        sid = await sm.create_session("prev task")
        await sm.add_message("user", "prev msg")

        class FirstStepFailsLoop(FakeReActLoop):
            def __init__(self):
                super().__init__()
                self.steps = []

            async def run(self, initial_prompt="", task_description="", todo_manager=None, history_messages=None, on_new_messages=None, session_id=None):
                self.last_messages = [
                    {"role": "system", "content": "sys"},
                    {"role": "user", "content": "failing task"},
                ]
                if history_messages:
                    self.last_messages.extend(history_messages)
                raise ValueError("API call failed immediately")

        agent = _build_agent(sm, td, FirstStepFailsLoop())
        result = await agent.execute(
            task="failing task",
            init_instructions="test",
            max_iterations=10,
        )

        assert result["status"] == "error"
        assert "API call failed" in result["error"]

        all_msgs = sm.current_messages
        contents = [m.content for m in all_msgs]
        assert "prev msg" in contents
        assert "failing task" in contents


@pytest.mark.asyncio
async def test_cancelled_error_propagates():
    """CancelledError 仍向上传播"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        await sm.create_session("test")

        agent = _build_agent(sm, td, FakeReActLoop(fail_mode="cancelled"))
        with pytest.raises(asyncio.CancelledError):
            await agent.execute(
                task="cancel me",
                init_instructions="test",
                max_iterations=10,
            )

        all_msgs = sm.current_messages
        contents = [m.content for m in all_msgs]
        # 现在用户消息在 react_loop 之前已保存，取消时也会保留
        assert "cancel me" in contents
        # 但 assistant 回复不应保存
        assert not any("✅ Task completed" in c for c in contents)


@pytest.mark.asyncio
async def test_tool_messages_saved_on_exception():
    """异常时 tool 消息仍保存（单轨 messages 中）"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        await sm.create_session("prev")

        agent = _build_agent(sm, td, FakeReActLoop(fail_mode="exception"))
        await agent.execute(
            task="error task",
            init_instructions="test",
            max_iterations=10,
        )

        f = sm.sessions_dir / f"{sm.current_session_id}.json"
        with open(f) as fh:
            raw = json.load(fh)
        assert "messages" in raw
        tool_msgs = [m for m in raw["messages"] if m.get("role") == "tool"]
        assert len(tool_msgs) >= 1


@pytest.mark.asyncio
async def test_reload_after_exception_preserves_data():
    """异常后重载 session，数据仍完整"""
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td)
        sm = SessionManager(project_path)
        sid = await sm.create_session("prev")
        await sm.add_message("user", "original message")

        agent = _build_agent(sm, td, FakeReActLoop(fail_mode="exception"))
        await agent.execute(
            task="error task",
            init_instructions="test",
            max_iterations=10,
        )

        sm2 = SessionManager(project_path)
        await sm2.switch_session(sid)
        msgs2 = await sm2.get_messages(include_system=False)
        c2 = [m["content"] for m in msgs2]
        assert "error task" in c2

        raw = json.loads(Path(sm.sessions_dir / f"{sid}.json").read_text())
        raw_msgs = [m["content"] for m in raw["messages"]]
        assert "original message" in raw_msgs
        assert "error task" in raw_msgs
