"""Tests for #1227 — structured message preservation in bootstrap/ingest/after_turn.

Verifies that multipart tool calls/results are stored as message_parts
and can be round-tripped through the conversation store.
"""

from __future__ import annotations

import json
import pytest
import pytest_asyncio

from cortex.hippo.bootstrap import extract_message_parts, flatten_content
from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine


# ═══════════════════════════════════════════════════════════════════════════
# Unit tests: extract_message_parts
# ═══════════════════════════════════════════════════════════════════════════


class TestExtractMessageParts:
    """Test part extraction from raw message dicts."""

    def test_plain_text_returns_empty(self):
        msg = {"role": "user", "content": "hello"}
        assert extract_message_parts(msg) == []

    def test_anthropic_tool_use(self):
        msg = {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "Let me check that."},
                {
                    "type": "tool_use",
                    "id": "call_123",
                    "name": "web_search",
                    "input": {"query": "test"},
                },
            ],
        }
        parts = extract_message_parts(msg)
        assert len(parts) == 2
        assert parts[0]["part_type"] == "text"
        assert parts[0]["text_content"] == "Let me check that."
        assert parts[1]["part_type"] == "tool_use"
        assert parts[1]["tool_call_id"] == "call_123"
        assert parts[1]["tool_name"] == "web_search"
        assert json.loads(parts[1]["tool_input"]) == {"query": "test"}

    def test_anthropic_tool_result(self):
        msg = {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": "call_123",
                    "content": "Search results: ...",
                },
            ],
        }
        parts = extract_message_parts(msg)
        assert len(parts) == 1
        assert parts[0]["part_type"] == "tool_result"
        assert parts[0]["tool_call_id"] == "call_123"
        assert parts[0]["tool_output"] == "Search results: ..."

    def test_openai_tool_calls(self):
        msg = {
            "role": "assistant",
            "content": "I'll search for that.",
            "tool_calls": [
                {
                    "id": "call_456",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": '{"query": "test"}',
                    },
                }
            ],
        }
        parts = extract_message_parts(msg)
        assert len(parts) == 2  # text + tool_call
        assert parts[0]["part_type"] == "text"
        assert parts[0]["text_content"] == "I'll search for that."
        assert parts[1]["part_type"] == "tool_use"
        assert parts[1]["tool_call_id"] == "call_456"
        assert parts[1]["tool_name"] == "web_search"

    def test_openai_tool_result(self):
        msg = {
            "role": "tool",
            "tool_call_id": "call_456",
            "name": "web_search",
            "content": "Result text",
        }
        parts = extract_message_parts(msg)
        assert len(parts) == 1
        assert parts[0]["part_type"] == "tool_result"
        assert parts[0]["tool_call_id"] == "call_456"
        assert parts[0]["tool_name"] == "web_search"
        assert parts[0]["tool_output"] == "Result text"

    def test_thinking_block(self):
        msg = {
            "role": "assistant",
            "content": [
                {"type": "thinking", "thinking": "Let me reason..."},
                {"type": "text", "text": "Here's my answer."},
            ],
        }
        parts = extract_message_parts(msg)
        assert len(parts) == 2
        assert parts[0]["part_type"] == "reasoning"
        assert parts[0]["text_content"] == "Let me reason..."
        assert parts[1]["part_type"] == "text"

    def test_mixed_tool_conversation(self):
        """Simulate a realistic multi-tool conversation."""
        messages = [
            {"role": "user", "content": "Search the web and read a file"},
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "I'll do both."},
                    {"type": "tool_use", "id": "c1", "name": "web_search", "input": {"q": "test"}},
                    {"type": "tool_use", "id": "c2", "name": "read_file", "input": {"path": "/tmp/x"}},
                ],
            },
            {
                "role": "user",
                "content": [
                    {"type": "tool_result", "tool_use_id": "c1", "content": "web result"},
                    {"type": "tool_result", "tool_use_id": "c2", "content": "file content"},
                ],
            },
        ]
        # First message: plain text, no parts
        assert extract_message_parts(messages[0]) == []
        # Second: 1 text + 2 tool_use
        parts1 = extract_message_parts(messages[1])
        assert len(parts1) == 3
        assert parts1[0]["part_type"] == "text"
        assert parts1[1]["tool_call_id"] == "c1"
        assert parts1[2]["tool_call_id"] == "c2"
        # Third: 2 tool_result
        parts2 = extract_message_parts(messages[2])
        assert len(parts2) == 2
        assert parts2[0]["tool_call_id"] == "c1"
        assert parts2[1]["tool_call_id"] == "c2"


# ═══════════════════════════════════════════════════════════════════════════
# Unit tests: flatten_content
# ═══════════════════════════════════════════════════════════════════════════


class TestFlattenContent:
    def test_string_passthrough(self):
        assert flatten_content({"content": "hello"}) == "hello"

    def test_empty(self):
        assert flatten_content({}) == ""

    def test_list_content(self):
        msg = {
            "content": [
                {"type": "text", "text": "Hello"},
                {"type": "tool_use", "name": "search", "input": {}},
            ]
        }
        result = flatten_content(msg)
        assert "Hello" in result
        assert "[tool_call: search]" in result

    def test_tool_result_flattened(self):
        msg = {
            "content": [
                {"type": "tool_result", "content": "result data"},
            ]
        }
        assert flatten_content(msg) == "result data"


# ═══════════════════════════════════════════════════════════════════════════
# Integration: round-trip through engine
# ═══════════════════════════════════════════════════════════════════════════


@pytest_asyncio.fixture
async def engine():
    config = HippoConfig(enabled=True, db_path=":memory:")
    eng = HippoEngine(config=config)
    await eng.initialize(db_path=":memory:", owner_id="test", conversation_id="test-sess")
    yield eng
    await eng.shutdown()


class TestStructuredIngestRoundTrip:
    """Verify structured messages survive ingest → store → retrieve."""

    @pytest.mark.asyncio
    async def test_ingest_message_with_parts(self, engine: HippoEngine):
        """ingest_message() stores parts when provided."""
        parts = [
            {"ordinal": 0, "part_type": "text", "text_content": "Let me search."},
            {
                "ordinal": 1,
                "part_type": "tool_use",
                "tool_call_id": "call_abc",
                "tool_name": "web_search",
                "tool_input": '{"query": "test"}',
            },
        ]
        result = await engine.ingest_message(
            role="assistant",
            content="Let me search.\n[tool_call: web_search]",
            message_parts=parts,
        )
        assert result.message_id is not None

        # Retrieve parts
        stored_parts = await engine.conversation_store.get_message_parts(result.message_id)
        assert len(stored_parts) == 2
        assert stored_parts[0]["part_type"] == "text"
        assert stored_parts[0]["text_content"] == "Let me search."
        assert stored_parts[1]["part_type"] == "tool_use"
        assert stored_parts[1]["tool_call_id"] == "call_abc"
        assert stored_parts[1]["tool_name"] == "web_search"

    @pytest.mark.asyncio
    async def test_bootstrap_preserves_tool_structure(self, engine: HippoEngine, tmp_path):
        """bootstrap() stores message_parts for structured messages."""
        session_data = [
            {"role": "user", "content": "Do a search"},
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "Searching now."},
                    {"type": "tool_use", "id": "tc_1", "name": "search", "input": {"q": "test"}},
                ],
            },
            {
                "role": "user",
                "content": [
                    {"type": "tool_result", "tool_use_id": "tc_1", "content": "Found it!"},
                ],
            },
            {"role": "assistant", "content": "Here are the results."},
        ]
        session_file = tmp_path / "session.json"
        session_file.write_text(json.dumps(session_data))

        result = await engine.bootstrap(
            session_id="bootstrap-test",
            session_file=str(session_file),
        )
        assert result["bootstrapped"] is True
        assert result["message_count"] == 4

        # Retrieve all messages and check parts
        conv = await engine.conversation_store.get_or_create_conversation(
            "test", "bootstrap-test"
        )
        messages = await engine.conversation_store.get_messages(
            conv.conversation_id, "test"
        )
        assert len(messages) >= 4

        # The assistant message (index 1) should have parts
        assistant_msg = messages[1]
        parts = await engine.conversation_store.get_message_parts(assistant_msg.message_id)
        assert len(parts) == 2
        assert parts[0]["part_type"] == "text"
        assert parts[0]["text_content"] == "Searching now."
        assert parts[1]["part_type"] == "tool_use"
        assert parts[1]["tool_call_id"] == "tc_1"
        assert parts[1]["tool_name"] == "search"

        # The tool result message (index 2) should have parts
        tool_result_msg = messages[2]
        tool_parts = await engine.conversation_store.get_message_parts(tool_result_msg.message_id)
        assert len(tool_parts) == 1
        assert tool_parts[0]["tool_call_id"] == "tc_1"
        assert tool_parts[0]["tool_output"] == "Found it!"

        # Plain text messages get text parts (our convention — full reconstruction)
        plain_user = messages[0]
        plain_user_parts = await engine.conversation_store.get_message_parts(plain_user.message_id)
        assert len(plain_user_parts) == 1
        assert plain_user_parts[0]["part_type"] == "text"
        plain_assistant = messages[3]
        plain_asst_parts = await engine.conversation_store.get_message_parts(plain_assistant.message_id)
        assert len(plain_asst_parts) == 1
        assert plain_asst_parts[0]["part_type"] == "text"

    @pytest.mark.asyncio
    async def test_after_turn_preserves_tool_structure(self, engine: HippoEngine):
        """after_turn() stores message_parts for structured messages."""
        messages = [
            {"role": "user", "content": "Use a tool"},
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "Using tool."},
                    {"type": "tool_use", "id": "tc_at", "name": "exec", "input": {"cmd": "ls"}},
                ],
            },
            {
                "role": "user",
                "content": [
                    {"type": "tool_result", "tool_use_id": "tc_at", "content": "file.txt"},
                ],
            },
            {"role": "assistant", "content": "Done!"},
        ]

        conv = engine._conversation
        await engine.after_turn(
            conversation_id=conv.conversation_id,
            messages=messages,
            pre_prompt_count=0,
        )

        # Fetch stored messages
        stored = await engine.conversation_store.get_messages(
            conv.conversation_id, "test"
        )
        assert len(stored) == 4

        # Check assistant message has tool parts
        assistant_parts = await engine.conversation_store.get_message_parts(
            stored[1].message_id
        )
        assert len(assistant_parts) == 2
        assert assistant_parts[1]["tool_call_id"] == "tc_at"
        assert assistant_parts[1]["tool_name"] == "exec"

        # Check tool result
        result_parts = await engine.conversation_store.get_message_parts(
            stored[2].message_id
        )
        assert len(result_parts) == 1
        assert result_parts[0]["tool_call_id"] == "tc_at"
        assert result_parts[0]["tool_output"] == "file.txt"
