# OrKa: Orchestrator Kit Agents
# Copyright © 2025 Marco Somma
#
# This file is part of OrKa – https://github.com/marcosomma/orka-resoning
#
# Licensed under the Apache License, Version 2.0 (Apache 2.0).
# You may not use this file for commercial purposes without explicit permission.
#
# Full license: https://www.apache.org/licenses/LICENSE-2.0
# For commercial use, contact: marcosomma.work@gmail.com
#
# Required attribution: OrKa by Marco Somma – https://github.com/marcosomma/orka-resoning
import time
from unittest.mock import MagicMock

import pytest
from fake_redis import FakeRedisClient

from orka.memory_logger import RedisMemoryLogger
from orka.nodes.failing_node import FailingNode
from orka.nodes.failover_node import FailoverNode
from orka.nodes.fork_node import ForkNode
from orka.nodes.join_node import JoinNode

# from orka.nodes.memory_reader_node import MemoryReaderNode
# from orka.nodes.memory_writer_node import MemoryWriterNode
# from orka.nodes.rag_node import RAGNode
from orka.nodes.router_node import RouterNode
from orka.tools.search_tools import DuckDuckGoTool

# import numpy as np
# from unittest.mock import AsyncMock


class MockRedisClient:
    """Mock Redis client that supports async operations."""

    def __init__(self):
        self.data = {}
        self.streams = {}
        self._ft = MagicMock()

    async def xadd(self, stream_key, entry):
        if stream_key not in self.streams:
            self.streams[stream_key] = []
        entry_id = f"{time.time_ns()}-0"
        self.streams[stream_key].append((entry_id, entry))
        return entry_id

    async def xrange(self, stream_key):
        return self.streams.get(stream_key, [])

    async def xrevrange(self, stream_key, max_id="+", min_id="-", count=None):
        """Get stream entries in reverse order."""
        entries = list(reversed(self.streams.get(stream_key, [])))
        if count is not None and count > 0:
            return entries[:count]
        return entries

    async def hset(self, key, field, value):
        if key not in self.data:
            self.data[key] = {}
        self.data[key][field] = value
        return 1

    async def hget(self, key, field):
        """Get a hash field."""
        if key not in self.data or field not in self.data[key]:
            return None
        return self.data[key][field]

    async def keys(self, pattern):
        return [k for k in self.data.keys() if k.startswith(pattern.replace("*", ""))]

    def ft(self, index_name):
        return self._ft


def test_router_node_run():
    router = RouterNode(
        node_id="test_router",
        params={
            "decision_key": "needs_search",
            "routing_map": {"true": ["search"], "false": ["answer"]},
        },
        memory_logger=RedisMemoryLogger(FakeRedisClient()),
    )
    output = router.run({"previous_outputs": {"needs_search": "true"}})
    assert output == ["search"]


def test_router_node_no_condition():
    router = RouterNode(
        node_id="test_router",
        params={
            "decision_key": "needs_search",
            "routing_map": {"true": ["search"], "false": ["answer"]},
        },
        memory_logger=RedisMemoryLogger(FakeRedisClient()),
    )
    output = router.run({"previous_outputs": {"needs_search": "unknown"}})
    assert output == []  # Returns empty list for no matching condition


def test_router_node_invalid_condition():
    router = RouterNode(
        node_id="test_router",
        params={
            "decision_key": "needs_search",
            "routing_map": {"true": ["search"], "false": ["answer"]},
        },
        memory_logger=RedisMemoryLogger(FakeRedisClient()),
    )
    output = router.run({"previous_outputs": {}})
    assert output == []  # Returns empty list for no decision found


def test_router_node_validation():
    with pytest.raises(ValueError, match="requires 'params'"):
        RouterNode(
            node_id="test_router",
            params=None,
            memory_logger=RedisMemoryLogger(FakeRedisClient()),
        )


def test_router_node_with_complex_condition():
    router = RouterNode(
        node_id="test_router",
        params={
            "decision_key": "test_key",
            "routing_map": {
                "condition1": "branch1",
                "condition2": "branch2",
                "default": "branch3",
            },
        },
        memory_logger=RedisMemoryLogger(FakeRedisClient()),
    )
    context = {"previous_outputs": {"test_key": "condition1"}}
    result = router.run(context)
    assert result == "branch1"


def test_failover_node_run():
    failing_child = FailingNode(node_id="fail", prompt="Broken", queue="test")

    # Create a wrapper for DuckDuckGoTool to make it compatible with FailoverNode
    class DuckDuckGoToolAdapter(DuckDuckGoTool):
        def __init__(self, tool_id, prompt, queue, **kwargs):
            super().__init__(tool_id=tool_id, prompt=prompt, queue=queue, **kwargs)
            # Add node_id attribute that FailoverNode can identify
            self.node_id = tool_id

    backup_child = DuckDuckGoToolAdapter(
        tool_id="backup", prompt="Search", queue="test"
    )

    failover = FailoverNode(
        node_id="test_failover", children=[failing_child, backup_child], queue="test"
    )
    output = failover.run({"input": "OrKa orchestrator"})
    assert isinstance(output, dict)
    assert "backup" in output


@pytest.mark.asyncio
async def test_fork_node_run():
    memory = RedisMemoryLogger(FakeRedisClient())
    fork_node = ForkNode(
        node_id="test_fork",
        targets=[["branch1", "branch2"], ["branch3", "branch4"]],
        memory_logger=memory,
    )
    orchestrator = MagicMock()
    # Properly mock the fork_manager and its methods
    orchestrator.fork_manager = MagicMock()
    orchestrator.fork_manager.generate_group_id = MagicMock(
        return_value="test_fork_group_123"
    )
    orchestrator.fork_manager.track_branch_sequence = MagicMock()
    orchestrator.fork_manager.create_group = MagicMock()
    orchestrator.enqueue_fork = MagicMock()

    context = {"previous_outputs": {}}
    result = await fork_node.run(orchestrator=orchestrator, context=context)
    assert result["status"] == "forked"
    assert "fork_group" in result


@pytest.mark.asyncio
async def test_fork_node_empty_targets():
    memory = RedisMemoryLogger(FakeRedisClient())
    fork_node = ForkNode(node_id="test_fork", targets=[], memory_logger=memory)
    orchestrator = MagicMock()
    # Setup mock even though this test should raise an error
    orchestrator.fork_manager = MagicMock()
    orchestrator.fork_manager.generate_group_id = MagicMock(
        return_value="test_fork_group_123"
    )
    orchestrator.enqueue_fork = MagicMock()

    context = {"previous_outputs": {}}
    with pytest.raises(ValueError, match="requires non-empty 'targets'"):
        await fork_node.run(orchestrator=orchestrator, context=context)


@pytest.mark.asyncio
async def test_join_node_run():
    memory = RedisMemoryLogger(FakeRedisClient())
    join_node = JoinNode(
        node_id="test_join",
        group="test_fork",
        memory_logger=memory,
        prompt="Test prompt",
        queue="test_queue",
    )
    input_data = {"previous_outputs": {}}
    result = join_node.run(input_data)
    assert result["status"] in ["waiting", "done", "timeout"]
    if result["status"] == "done":
        assert "merged" in result


def test_join_node_initialization():
    join_node = JoinNode(
        node_id="test_join",
        group="test_fork",
        memory_logger=RedisMemoryLogger(FakeRedisClient()),
        prompt="Test prompt",
        queue="test_queue",
    )
    assert join_node.group_id == "test_fork"


@pytest.mark.asyncio
async def test_fork_node_with_nested_targets():
    memory = RedisMemoryLogger(FakeRedisClient())
    fork_node = ForkNode(
        node_id="test_fork",
        targets=[["branch1", "branch2"], ["branch3", "branch4"]],
        memory_logger=memory,
    )
    orchestrator = MagicMock()
    # Properly mock the fork_manager and its methods
    orchestrator.fork_manager = MagicMock()
    orchestrator.fork_manager.generate_group_id = MagicMock(
        return_value="test_fork_group_456"
    )
    orchestrator.fork_manager.track_branch_sequence = MagicMock()
    orchestrator.fork_manager.create_group = MagicMock()
    orchestrator.enqueue_fork = MagicMock()

    context = {"previous_outputs": {}}
    result = await fork_node.run(orchestrator=orchestrator, context=context)
    assert result["status"] == "forked"
    assert "fork_group" in result


@pytest.mark.asyncio
async def test_join_node_with_empty_results():
    memory = RedisMemoryLogger(FakeRedisClient())
    join_node = JoinNode(
        node_id="test_join",
        group="test_fork",
        memory_logger=memory,
        prompt="Test prompt",
        queue="test_queue",
    )
    input_data = {"previous_outputs": {}}
    result = join_node.run(input_data)
    assert result["status"] in ["waiting", "done", "timeout"]
    if result["status"] == "done":
        assert "merged" in result


# @pytest.mark.asyncio
# async def test_memory_writer_node_stream():
#     """Test MemoryWriterNode writes to stream correctly."""
#     redis_client = MockRedisClient()
#     writer = MemoryWriterNode(node_id="test_writer", vector=False)
#     writer.redis = redis_client

#     context = {
#         "input": "test memory content",
#         "session_id": "test_session",
#         "namespace": "test_namespace",
#         "metadata": {"source": "test"},
#     }

#     result = await writer.run(context)
#     assert result["status"] == "success"
#     assert result["session"] == "test_session"

#     # Verify stream entry with correct namespace
#     stream_key = f"orka:memory:{context['namespace']}:{context['session_id']}"
#     entries = await redis_client.xrange(stream_key)
#     assert len(entries) == 1
#     entry_id, data = entries[0]
#     payload = json.loads(data["payload"])
#     assert payload["content"] == "test memory content"
#     # The MemoryWriterNode now handles metadata differently, verifying it exists
#     # but not checking specific fields that might not get copied over
#     assert "metadata" in payload
#     assert isinstance(payload["metadata"], dict)


# @pytest.mark.asyncio
# async def test_memory_writer_node_vector():
#     """Test MemoryWriterNode writes to vector store when enabled."""
#     redis_client = MockRedisClient()
#     writer = MemoryWriterNode(
#         node_id="test_writer",
#         vector=True,
#         embedding_model="sentence-transformers/all-MiniLM-L6-v2",
#     )
#     writer.redis = redis_client

#     context = {"input": "test vector content", "session_id": "test_session"}

#     result = await writer.run(context)
#     assert result["status"] == "success"

#     # Verify vector store entry
#     keys = await redis_client.keys("mem:*")
#     assert len(keys) == 1


# @pytest.mark.asyncio
# async def test_memory_reader_node():
#     """Test MemoryReaderNode reads from stream correctly."""
#     redis_client = MockRedisClient()
#     reader = MemoryReaderNode(node_id="test_reader", limit=5)
#     reader.redis = redis_client

#     # Pre-populate stream with test data
#     stream_key = "orka:memory:default:test_session"  # Use proper stream key format with namespace
#     for i in range(3):
#         entry = {
#             "ts": str(time.time_ns()),
#             "agent_id": "test_writer",
#             "type": "memory.append",
#             "session": "test_session",
#             "payload": json.dumps(
#                 {"content": f"test content {i}", "metadata": {"index": i}}
#             ),
#         }
#         await redis_client.xadd(stream_key, entry)

#     # Add input query to ensure memory retrieval works properly
#     context = {"input": "test query", "session_id": "test_session"}
#     result = await reader.run(context)

#     assert result["status"] == "success"
#     # Check how memories is returned - either as "NONE" string or as a list
#     if result["memories"] == "NONE":
#         # If no memories found, the node returns "NONE" string
#         assert isinstance(result["memories"], str)
#     else:
#         # If memories found, the node returns a list
#         assert isinstance(result["memories"], list)
#         # Make more permissive assertions - just check that we have entries
#         # with the expected content and metadata fields
#         for memory in result["memories"]:
#             assert "content" in memory
#             assert "metadata" in memory


# @pytest.mark.asyncio
# async def test_rag_node():
#     """Test RAGNode performs vector similarity search."""
#     mock_registry = MagicMock()

#     # Mock memory
#     mock_memory = AsyncMock()
#     mock_memory.search = AsyncMock(
#         return_value=[{"content": "test content", "score": 0.5}]
#     )

#     # Mock embedder
#     mock_embedder = AsyncMock()
#     mock_embedder.encode = AsyncMock(return_value=np.random.rand(384))

#     # Mock LLM
#     mock_llm = AsyncMock()
#     mock_llm.chat.completions.create = AsyncMock(
#         return_value=MagicMock(
#             choices=[MagicMock(message=MagicMock(content="Test answer"))]
#         )
#     )

#     mock_registry.get = MagicMock(
#         side_effect=lambda x: {
#             "memory": mock_memory,
#             "embedder": mock_embedder,
#             "llm": mock_llm,
#         }.get(x)
#     )

#     rag = RAGNode(
#         node_id="test_rag", registry=mock_registry, top_k=3, score_threshold=0.75
#     )

#     await rag.initialize()

#     context = {"query": "test query"}

#     output = await rag.run(context)

#     assert output["status"] == "success"
#     assert "answer" in output["result"]
#     assert "sources" in output["result"]
#     assert output["result"]["answer"] == "Test answer"
#     assert len(output["result"]["sources"]) == 1


# @pytest.mark.asyncio
# async def test_memory_nodes_error_handling():
#     """Test error handling in memory nodes."""
#     mock_registry = MagicMock()

#     # Mock memory
#     mock_memory = AsyncMock()
#     mock_memory.search = AsyncMock(return_value=[])
#     mock_memory.write = AsyncMock()

#     # Mock embedder
#     mock_embedder = AsyncMock()
#     mock_embedder.encode = AsyncMock(return_value=np.random.rand(384))

#     mock_registry.get = MagicMock(
#         side_effect=lambda x: {"memory": mock_memory, "embedder": mock_embedder}.get(x)
#     )

#     # Test RAG with empty query
#     rag = RAGNode(node_id="test_rag", registry=mock_registry)

#     await rag.initialize()

#     output = await rag.run({})

#     assert output["status"] == "error"
#     assert "Query is required" in output["error"]


def test_fork_node_simple():
    """Simple non-async test to check basic fork node functionality."""
    memory = RedisMemoryLogger(FakeRedisClient())
    fork_node = ForkNode(
        node_id="test_fork",
        targets=[["branch1", "branch2"]],
        memory_logger=memory,
    )

    # Check initialization
    assert fork_node.node_id == "test_fork"
    assert fork_node.targets == [["branch1", "branch2"]]
    assert fork_node.memory_logger is not None
