"""
Tests for ApeChain - High-level wrapper for Ape task execution.
"""

import pytest
from unittest.mock import Mock, patch, MagicMock

from ape_langchain.chain import ApeChain
from ape_langchain.core import ApeModule, ApeExecutionError


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_initialization(mock_validate, mock_compile):
    """Test ApeChain initialization."""
    # Mock module and signature
    mock_signature = Mock()
    mock_signature.name = "process"
    mock_signature.inputs = {"data": "str"}
    mock_signature.output = "str"
    mock_signature.description = "Process data"
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["process"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain
    chain = ApeChain(ape_file="test.ape", function="process")
    
    assert chain.ape_file == "test.ape"
    assert chain.function == "process"
    
    mock_compile.assert_called_once_with("test.ape")
    mock_validate.assert_called_once_with(mock_module)


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_from_ape_file(mock_validate, mock_compile):
    """Test ApeChain.from_ape_file() class method."""
    # Mock module and signature
    mock_signature = Mock()
    mock_signature.name = "transform"
    mock_signature.inputs = {"input": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["transform"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain using class method
    chain = ApeChain.from_ape_file("test.ape", function="transform")
    
    assert isinstance(chain, ApeChain)
    assert chain.ape_file == "test.ape"
    assert chain.function == "transform"


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
@patch("ape_langchain.chain.execute_ape_function")
def test_ape_chain_run_success(mock_execute, mock_validate, mock_compile):
    """Test successful ApeChain execution."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "double"
    mock_signature.inputs = {"value": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["double"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    mock_execute.return_value = "10"
    
    # Create and run chain
    chain = ApeChain.from_ape_file("test.ape", function="double")
    result = chain.run(value="5")
    
    assert result == "10"
    mock_execute.assert_called_once_with(mock_module, "double", value="5")


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_run_validation_error(mock_validate, mock_compile):
    """Test ApeChain execution with validation error."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "process"
    mock_signature.inputs = {"required_param": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["process"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain
    chain = ApeChain.from_ape_file("test.ape", function="process")
    
    # Run with missing parameter - should raise TypeError
    with pytest.raises(TypeError, match="Missing required parameters"):
        chain.run(wrong_param="value")


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
@patch("ape_langchain.chain.execute_ape_function")
def test_ape_chain_run_execution_error(mock_execute, mock_validate, mock_compile):
    """Test ApeChain execution with runtime error."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "fail"
    mock_signature.inputs = {"input": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["fail"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    mock_execute.side_effect = ApeExecutionError("Runtime error")
    
    # Create chain
    chain = ApeChain.from_ape_file("test.ape", function="fail")
    
    # Run should raise ApeExecutionError
    with pytest.raises(ApeExecutionError, match="Runtime error"):
        chain.run(input="test")


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_get_task(mock_validate, mock_compile):
    """Test ApeChain.get_task() method."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "calculate"
    mock_signature.inputs = {"x": "str", "y": "str"}
    mock_signature.output = "str"
    mock_signature.description = "Calculate something"
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["calculate"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain and get task
    chain = ApeChain.from_ape_file("test.ape", function="calculate")
    task = chain.get_task()
    
    assert task.name == "calculate"
    assert task.inputs == {"x": "str", "y": "str"}
    assert task.output == "str"


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_get_json_schema(mock_validate, mock_compile):
    """Test ApeChain.get_json_schema() method."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "transform"
    mock_signature.inputs = {"data": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["transform"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain and get schema
    chain = ApeChain.from_ape_file("test.ape", function="transform")
    schema = chain.get_json_schema()
    
    assert schema["type"] == "function"
    assert schema["function"]["name"] == "transform"
    assert "data" in schema["function"]["parameters"]["properties"]


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
@patch("ape_langchain.chain.execute_ape_function")
def test_ape_chain_callable(mock_execute, mock_validate, mock_compile):
    """Test that ApeChain can be called directly like a function."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "add"
    mock_signature.inputs = {"a": "str", "b": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["add"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    mock_execute.return_value = "8"
    
    # Create chain
    chain = ApeChain.from_ape_file("test.ape", function="add")
    
    # Call directly
    result = chain(a="5", b="3")
    
    assert result == "8"
    mock_execute.assert_called_once_with(mock_module, "add", a="5", b="3")


@patch("ape_langchain.chain.compile_ape")
@patch("ape_langchain.chain.validate_ape")
def test_ape_chain_repr(mock_validate, mock_compile):
    """Test ApeChain string representation."""
    # Setup mocks
    mock_signature = Mock()
    mock_signature.name = "process"
    mock_signature.inputs = {"input": "str"}
    mock_signature.output = "str"
    mock_signature.description = None
    
    mock_module = Mock(spec=ApeModule)
    mock_module.get_function_signature.return_value = mock_signature
    mock_module.list_functions.return_value = ["process"]
    
    mock_compile.return_value = mock_module
    mock_validate.return_value = None
    
    # Create chain
    chain = ApeChain.from_ape_file("test.ape", function="process")
    repr_str = repr(chain)
    
    assert "ApeChain" in repr_str
    assert "test.ape" in repr_str
    assert "process" in repr_str
