import json

import numpy as np

from masfactory.adapters.context.types import ContextQuery
from masfactory.adapters.memory import HistoryMemory, VectorMemory
from masfactory.adapters.model import ModelResponseType
from masfactory.adapters.retrieval import FileSystemRetriever, SimpleKeywordRetriever, VectorRetriever
from masfactory.components.agents.agent import Agent
from masfactory.core.message import ParagraphMessageFormatter


def _embed_counts(text: str) -> np.ndarray:
    lowered = (text or "").lower()
    # Deterministic small embedding for unit tests.
    return np.array(
        [
            lowered.count("coffee"),
            lowered.count("tea"),
            lowered.count("banana"),
        ],
        dtype=float,
    )


def test_vector_retriever_get_blocks_selects_relevant_doc():
    retriever = VectorRetriever(
        documents={"a": "I like coffee", "b": "tea only", "c": "banana split"},
        embedding_function=_embed_counts,
        similarity_threshold=0.5,
    )
    blocks = retriever.get_blocks(ContextQuery(query_text="coffee"), top_k=1)
    assert len(blocks) == 1
    assert blocks[0].text == "I like coffee"


def test_filesystem_retriever_get_blocks_selects_relevant_file(tmp_path):
    docs_dir = tmp_path / "docs"
    docs_dir.mkdir()
    (docs_dir / "a.txt").write_text("coffee coffee", encoding="utf-8")
    (docs_dir / "b.txt").write_text("tea", encoding="utf-8")

    retriever = FileSystemRetriever(
        docs_dir=docs_dir,
        embedding_function=_embed_counts,
        similarity_threshold=0.5,
        file_extension=".txt",
    )
    blocks = retriever.get_blocks(ContextQuery(query_text="coffee"), top_k=1)
    assert len(blocks) == 1
    assert blocks[0].text == "coffee coffee"
    assert blocks[0].metadata.get("doc_id") == "a.txt"


def test_simple_keyword_retriever_prefers_documents_with_query_terms():
    retriever = SimpleKeywordRetriever(
        documents={"a": "coffee and tea", "b": "tea only", "c": "banana split"},
    )
    blocks = retriever.get_blocks(ContextQuery(query_text="coffee"), top_k=1)
    assert len(blocks) == 1
    assert blocks[0].metadata.get("doc_id") == "a"


def test_agent_observe_injects_retriever_blocks_into_prompt():
    retriever = SimpleKeywordRetriever(documents={"a": "coffee and tea"}, passive=True, active=False)
    agent = Agent(
        name="a",
        instructions="do",
        model=object(),
        formatters=ParagraphMessageFormatter(),
        retrievers=[retriever],
    )
    _, user_prompt, _ = agent.observe({"input": "coffee"})
    assert "CONTEXT:" in user_prompt
    assert "(KEYWORD_RETRIEVER) coffee and tea" in user_prompt


class _FakeToolCallingModel:
    """Minimal model stub that triggers one tool call then returns content."""

    def __init__(self, *, tool_source: str, tool_query: str):
        self._tool_source = tool_source
        self._tool_query = tool_query
        self._calls = 0

    def invoke(self, *, messages, tools, settings=None, **kwargs):  # noqa: D401, ARG002
        self._calls += 1
        if self._calls == 1:
            return {
                "type": ModelResponseType.TOOL_CALL,
                "content": [
                    {
                        "id": "call_1",
                        "name": "retrieve_context",
                        "arguments": {"source": self._tool_source, "query": self._tool_query, "top_k": 1},
                    }
                ],
            }

        tool_messages = [m for m in (messages or []) if m.get("role") == "tool"]
        assert tool_messages, "Expected at least one tool message after tool call"
        payload = json.loads(tool_messages[-1]["content"])
        first_text = payload["blocks"][0]["text"]
        return {
            "type": ModelResponseType.CONTENT,
            "content": f"RESULT:\n{first_text}",
        }


def test_agent_step_tool_call_retrieve_context_end_to_end():
    retriever = SimpleKeywordRetriever(
        documents={"a": "coffee and tea", "b": "tea only"},
        passive=False,
        active=True,
    )
    agent = Agent(
        name="a",
        instructions="do",
        model=_FakeToolCallingModel(tool_source="KEYWORD_RETRIEVER", tool_query="coffee"),
        formatters=ParagraphMessageFormatter(),
        retrievers=[retriever],
    )
    out = agent.step({"input": "ignored"})
    assert out["RESULT"] == "coffee and tea"


class _FakeEchoModel:
    """Model stub that always returns a paragraph response containing the input text."""

    def invoke(self, *, messages, tools, settings=None, **kwargs):  # noqa: D401, ARG002
        user = next((m.get("content") for m in reversed(messages or []) if m.get("role") == "user"), "")
        return {"type": ModelResponseType.CONTENT, "content": f"ANSWER:\n{user}"}


def test_agent_step_inserts_into_vector_memory_and_is_injected_next_turn():
    mem = VectorMemory(
        embedding_function=_embed_counts,
        query_threshold=0.0,
        memory_size=10,
        passive=True,
        active=False,
    )
    agent = Agent(
        name="a",
        instructions="do",
        model=_FakeEchoModel(),
        formatters=ParagraphMessageFormatter(),
        memories=[mem],
    )

    # Step 1 writes into memory.
    agent.step({"input": "coffee"})

    # Step 2 should inject the stored memory snippet into the user prompt.
    _, user_prompt, _ = agent.observe({"input": "coffee"})
    assert "CONTEXT:" in user_prompt
    assert "coffee" in user_prompt


def test_agent_history_memory_roundtrip_messages_in_observe():
    history = HistoryMemory(top_k=2, memory_size=100)
    agent = Agent(
        name="a",
        instructions="do",
        model=_FakeEchoModel(),
        formatters=ParagraphMessageFormatter(),
        memories=[history],
    )

    agent.step({"input": "first"})
    _, _, messages = agent.observe({"input": "second"})

    # Observe should include history messages between system and user.
    roles = [m["role"] for m in messages]
    assert roles[0] == "system"
    assert roles[-1] == "user"
    assert "assistant" in roles and "user" in roles
