import pytest
import asyncio
from unittest.mock import MagicMock, AsyncMock
from orchestra import enhance
from orchestra.adapters.langchain import EnhancedLangChain

class MockRunnable:
    def __init__(self):
        self.invoke_count = 0
        self.ainvoke_count = 0
    
    def invoke(self, input, config=None, **kwargs):
        self.invoke_count += 1
        return f"processed {input}"

    async def ainvoke(self, input, config=None, **kwargs):
        self.ainvoke_count += 1
        return f"processed {input}"

@pytest.fixture
def mock_runnable():
    return MockRunnable()

def test_langchain_sync_caching(mock_runnable):
    enhanced = enhance(mock_runnable)
    
    # 1. First call - miss
    res1 = enhanced.invoke("hello")
    assert res1 == "processed hello"
    assert mock_runnable.invoke_count == 1
    
    # 2. Second call - hit
    res2 = enhanced.invoke("hello")
    assert res2 == "processed hello"
    assert mock_runnable.invoke_count == 1 # Count should not increase

    # 3. New query - miss
    res3 = enhanced.invoke("world")
    assert res3 == "processed world"
    assert mock_runnable.invoke_count == 2

@pytest.mark.asyncio
async def test_langchain_async_caching(mock_runnable):
    enhanced = enhance(mock_runnable)
    
    # 1. First call - miss
    res1 = await enhanced.ainvoke("hello")
    assert res1 == "processed hello"
    assert mock_runnable.ainvoke_count == 1
    
    # 2. Second call - hit
    res2 = await enhanced.ainvoke("hello")
    assert res2 == "processed hello"
    assert mock_runnable.ainvoke_count == 1 # Count should not increase

    # 3. New query - miss
    res3 = await enhanced.ainvoke("world")
    assert res3 == "processed world"
    assert mock_runnable.ainvoke_count == 2

def test_langchain_attribute_delegation(mock_runnable):
    enhanced = enhance(mock_runnable)
    mock_runnable.some_custom_attr = "custom"
    assert enhanced.some_custom_attr == "custom"
