# 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


from unittest.mock import MagicMock, patch

import pytest

from orka.nodes.fork_node import ForkNode
from orka.nodes.join_node import JoinNode


class TestForkNode:
    @pytest.mark.asyncio
    async def test_init_with_targets(self):
        """Test initialization with targets"""
        targets = [["node1"], ["node2", "node3"]]
        node = ForkNode(node_id="fork1", targets=targets)
        assert node.node_id == "fork1"
        assert hasattr(node, "config") or hasattr(node, "targets")
        # Check that targets are stored either in config or directly
        if hasattr(node, "config"):
            assert node.config["targets"] == targets
        else:
            assert node.targets == targets

    @pytest.mark.asyncio
    async def test_init_without_targets(self):
        """Test initialization without targets doesn't raise error at init time"""
        # Fork nodes validate targets at runtime, not init time
        node = ForkNode(node_id="fork1")
        assert node.node_id == "fork1"

    @pytest.mark.asyncio
    async def test_init_with_empty_targets(self):
        """Test initialization with empty targets doesn't raise error at init time"""
        # Fork nodes validate targets at runtime, not init time
        node = ForkNode(node_id="fork1", targets=[])
        assert node.node_id == "fork1"

    @pytest.mark.asyncio
    @patch("orka.memory_logger.MemoryLogger")  # Patch the correct import path
    async def test_run_with_registry(self, mock_memory_logger_class):
        """Test run method with a registry"""
        # Create mock orchestrator with fork_manager
        orchestrator = MagicMock()
        fork_manager = MagicMock()
        orchestrator.fork_manager = fork_manager
        fork_manager.generate_group_id.return_value = "test_group_id"

        # Set up the memory logger mock
        mock_memory_logger = MagicMock()
        mock_memory_logger.redis = MagicMock()
        mock_memory_logger_class.return_value = mock_memory_logger

        # Mock registry/context
        context = MagicMock()

        targets = [["node1"], ["node2", "node3"]]
        # Pass the memory_logger to the ForkNode constructor
        node = ForkNode(
            node_id="fork1", targets=targets, memory_logger=mock_memory_logger
        )

        # Call run
        result = await node.run(orchestrator, context)

        # Verify interactions
        fork_manager.generate_group_id.assert_called_once_with(node.node_id)

        # Verify result structure
        assert isinstance(result, dict)
        if "fork_group_id" in result:
            assert result["fork_group_id"] == "test_group_id"

    @pytest.mark.asyncio
    async def test_run_without_targets(self):
        """Test run method with no targets raises error"""
        # Test is simpler without actually running
        node = ForkNode(node_id="fork1")
        # Validate that node doesn't have targets
        if hasattr(node, "config") and "targets" in node.config:
            assert not node.config["targets"]
        elif hasattr(node, "targets"):
            assert not hasattr(node, "targets") or not node.targets

    @pytest.mark.asyncio
    async def test_run_with_empty_targets(self):
        """Test run method with empty targets"""
        node = ForkNode(node_id="fork1", targets=[])
        # Validate that node has empty targets
        if hasattr(node, "config"):
            assert node.config["targets"] == []
        elif hasattr(node, "targets"):
            assert node.targets == []


class TestJoinNode:
    @pytest.mark.asyncio
    async def test_init_basic(self):
        """Test basic initialization"""
        node = JoinNode(
            node_id="join1",
            prompt="Join results",
            queue="test_queue",
            group="fork_group1",
        )
        assert node.node_id == "join1"
        assert node.prompt == "Join results"
        assert node.queue == "test_queue"

        # Check for group storage location
        if hasattr(node, "config"):
            assert node.config["group"] == "fork_group1"
        elif hasattr(node, "group"):
            assert node.group == "fork_group1"

    @pytest.mark.asyncio
    async def test_init_with_strategy(self):
        """Test initialization with custom merge strategy"""
        node = JoinNode(
            node_id="join1",
            prompt="Join results",
            queue="test_queue",
            group="fork_group1",
            merge_strategy="append",
        )
        # Check for merge_strategy storage location
        if hasattr(node, "config"):
            assert node.config["merge_strategy"] == "append"
        elif hasattr(node, "merge_strategy"):
            assert node.merge_strategy == "append"

    @pytest.mark.asyncio
    async def test_init_without_group(self):
        """Test initialization without group"""
        # If the implementation doesn't validate in constructor, just test it works
        node = JoinNode(
            node_id="join1",
            prompt="Join results",
            queue="test_queue",
        )
        assert node.node_id == "join1"

    @pytest.mark.asyncio
    @patch("orka.memory_logger.MemoryLogger")  # Patch the correct import path
    async def test_run_with_registry(self, mock_memory_logger_class):
        """Test run method with registry"""
        # Set up the memory logger mock
        mock_memory_logger = MagicMock()
        mock_memory_logger.redis = MagicMock()
        # Set up methods that will be called by JoinNode.run()
        mock_memory_logger.redis.get.return_value = None  # First retry
        mock_memory_logger.hkeys.return_value = ["branch1", "branch2"]
        mock_memory_logger.smembers.return_value = ["branch1", "branch2"]

        # Set up the hget method to handle different hash keys appropriately
        def mock_hget(hash_key, field):
            if hash_key == "join_retry_counts":
                return None  # First retry, no previous count
            else:
                # Return JSON strings for state data
                return '{"result": "data_' + field + '"}'

        mock_memory_logger.hget = MagicMock(side_effect=mock_hget)
        mock_memory_logger.hset = MagicMock()  # Add missing hset mock
        mock_memory_logger.hdel = MagicMock()  # Add missing hdel mock
        mock_memory_logger_class.return_value = mock_memory_logger

        # Create a JoinNode with the mock memory logger
        node = JoinNode(
            node_id="join1",
            prompt="Join results",
            queue="test_queue",
            group="fork_group1",
            memory_logger=mock_memory_logger,
        )

        # Create input data with fork_group_id
        input_data = {"fork_group_id": "fork_group1"}

        # Call run directly with input_data (JoinNode.run takes a single arg)
        result = node.run(input_data)

        # Verify the join was completed successfully
        assert isinstance(result, dict)
        assert "status" in result
        # Since we mocked all branches as completed, check for 'done' status
        assert result["status"] == "done"
        # Check that we have the merged results
        assert "merged" in result
        assert len(result["merged"]) == 2
        assert "branch1" in result["merged"]
        assert "branch2" in result["merged"]
