"""
Tests for the AgentRunner module using real agent calls.

This test suite covers the AgentRunner class functionality including:
- Real agent initialization and calls
- Batch processing with actual responses
- Database integration with real data
- Error handling scenarios
"""

import asyncio
import os
import tempfile
import pytest
from typing import List, Optional, Any
import pandas as pd
import easier as ezr
from pydantic import BaseModel, Field

from easier.agent_runner import AgentRunner, cleanup_all_agents, cancel_all_running_tasks, DatabaseSchemaValidationError


@pytest.fixture(autouse=True)
def cleanup_before_each_test():
    """Automatically cleanup before each test to ensure clean state"""
    cancel_all_running_tasks()
    yield
    # Cleanup after test as well
    cleanup_all_agents()


@pytest.fixture
def real_agent():
    """Create a real EZAgent for testing with simple system prompt"""
    system_prompt = "You are a helpful assistant. Answer questions briefly in 1-2 words."
    agent = ezr.EZAgent(system_prompt)
    return agent


@pytest.fixture
def temp_db_file():
    """Create a temporary database file path for testing"""
    # Create a temp directory and generate a unique db filename
    import tempfile
    import uuid
    temp_dir = tempfile.gettempdir()
    db_file = os.path.join(temp_dir, f"test_agent_runner_{uuid.uuid4().hex}.db")
    yield db_file
    # Cleanup
    try:
        os.unlink(db_file)
    except FileNotFoundError:
        pass


@pytest.fixture
def simple_prompts():
    """Simple prompts for testing that won't use many tokens"""
    return [
        "What is 1+1?",
        "Color of sky?", 
        "Capital of USA?",
        "Say hello"
    ]


@pytest.fixture
def pydantic_prompts():
    """Prompts that work well with structured Pydantic responses"""
    return [
        "Explain why water boils at 100°C. Think through this step by step.",
        "What are the benefits of exercise? Consider multiple aspects.",
        "How does photosynthesis work? Break down the process."
    ]


@pytest.fixture
def simple_framer_func():
    """Simple framer function that converts results to DataFrame"""
    def framer(results: List[Optional[Any]]) -> pd.DataFrame:
        data = []
        for i, result in enumerate(results):
            if result is not None:
                data.append({
                    "prompt_index": i, 
                    "response": str(result.output),
                    "success": True
                })
            else:
                data.append({
                    "prompt_index": i, 
                    "response": "FAILED",
                    "success": False
                })
        return pd.DataFrame(data)
    return framer


@pytest.fixture 
def pydantic_response_model():
    """Pydantic model for structured agent responses"""
    class ResponseModel(BaseModel):
        answer: str = Field(description="The response generated by the agent")
        thoughts: List[str] = Field(description="List of thoughts or reasoning steps")
        
    return ResponseModel


@pytest.fixture
def pydantic_framer_func():
    """Framer function that uses Pydantic model_dump() to create DataFrame rows"""
    def framer(results: List[Optional[Any]]) -> pd.DataFrame:
        data = []
        for i, result in enumerate(results):
            if result is not None and hasattr(result, 'output') and hasattr(result.output, 'model_dump'):
                # Use model_dump() to convert Pydantic model to dict
                row_data = result.output.model_dump()
                row_data["prompt_index"] = i
                row_data["success"] = True
                data.append(row_data)
            else:
                data.append({
                    "prompt_index": i,
                    "answer": "FAILED",
                    "thoughts": [],
                    "success": False
                })
        return pd.DataFrame(data)
    return framer


class TestAgentRunnerInitialization:
    """Test AgentRunner initialization and validation"""
    
    def test_init_with_valid_agent(self, real_agent):
        """Test initialization with valid agent"""
        with AgentRunner(real_agent) as runner:
            assert runner.agent == real_agent
            assert runner.db_enabled is False
            assert runner.db_file is None
            assert runner.table_name == "results"
            assert runner.overwrite is False
            assert runner.timeout == 300.0  # Default timeout
            assert hasattr(runner, '_cost_model')
            assert runner._cost_model.input_ppm_cost > 0
            assert runner._cost_model.output_ppm_cost > 0
            assert runner._cost_model.thought_ppm_cost > 0
            assert len(runner.active_tasks) == 0
            assert runner._is_running is False
            assert runner._cleanup_done is False
    
    def test_init_with_invalid_agent(self):
        """Test initialization with invalid agent type"""
        with pytest.raises(TypeError, match="Expected agent to be of type ezr.EZAgent"):
            AgentRunner("not an agent")
    
    def test_init_with_database_config(self, real_agent, temp_db_file):
        """Test initialization with database configuration"""
        with AgentRunner(
            real_agent, 
            db_file=temp_db_file, 
            overwrite=True, 
            table_name="test_table",
            timeout=60.0
        ) as runner:
            assert runner.db_enabled is True
            assert runner.db_file == temp_db_file
            assert runner.table_name == "test_table"
            assert runner.overwrite is True
            assert runner.timeout == 60.0
            # Cost model should be auto-detected from agent
            assert hasattr(runner, '_cost_model')
            assert runner._cost_model.input_ppm_cost > 0
            assert runner._cost_model.output_ppm_cost > 0
            assert runner._cost_model.thought_ppm_cost > 0
    
    def test_cost_model_auto_detection(self, real_agent):
        """Test that cost model is auto-detected from agent.model_name"""
        with AgentRunner(real_agent) as runner:
            # Should auto-detect cost model from agent
            assert hasattr(runner, '_cost_model')
            
            # Should have selected a cost model from COST_CONFIG
            from easier.agent_runner import COST_CONFIG
            model_names = list(COST_CONFIG.models.keys())
            
            # The cost model should match one of the available models
            found_match = False
            for model_name, cost_model in COST_CONFIG.models.items():
                if (runner._cost_model.input_ppm_cost == cost_model.input_ppm_cost and
                    runner._cost_model.output_ppm_cost == cost_model.output_ppm_cost and
                    runner._cost_model.thought_ppm_cost == cost_model.thought_ppm_cost):
                    found_match = True
                    break
            
            assert found_match, f"Cost model doesn't match any available models: {model_names}"
    
    def test_cost_model_fallback_with_warning(self, real_agent):
        """Test cost model fallback when agent model not found"""
        import warnings
        
        # Temporarily patch the agent's model_name to an unknown model
        original_model_name = real_agent.model_name
        real_agent.model_name = "unknown-model"
        
        try:
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter("always")
                
                with AgentRunner(real_agent) as runner:
                    # Should have fallen back to flash model
                    from easier.agent_runner import COST_CONFIG
                    flash_model = COST_CONFIG.models["google-vertex:gemini-2.5-flash"]
                    
                    assert runner._cost_model.input_ppm_cost == flash_model.input_ppm_cost
                    assert runner._cost_model.output_ppm_cost == flash_model.output_ppm_cost
                    assert runner._cost_model.thought_ppm_cost == flash_model.thought_ppm_cost
                
                # Should have issued a warning
                assert len(w) == 1
                assert "Cost model 'unknown-model' not found" in str(w[0].message)
                assert "Falling back to 'google-vertex:gemini-2.5-flash'" in str(w[0].message)
        finally:
            # Restore original model_name
            real_agent.model_name = original_model_name

    def test_init_database_overwrite(self, real_agent, temp_db_file):
        """Test database initialization with overwrite"""
        # Create the file first
        with open(temp_db_file, 'w') as f:
            f.write("test")
        
        original_exists = os.path.exists(temp_db_file)
        assert original_exists
        
        with AgentRunner(real_agent, db_file=temp_db_file, overwrite=True) as runner:
            # File should have been removed - the AgentRunner only deletes it,
            # doesn't recreate until first write
            assert not os.path.exists(temp_db_file)
    
    def test_context_manager_cleanup(self, real_agent):
        """Test that context manager properly cleans up"""
        runner = AgentRunner(real_agent)
        
        with runner:
            assert runner._cleanup_done is False
        
        # After exiting context, cleanup should be done
        assert runner._cleanup_done is True


class TestBatchCreation:
    """Test batch creation functionality"""
    
    def test_create_batches_exact_division(self, real_agent):
        """Test batch creation when prompts divide evenly"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?", "Color of sky?", "Capital of USA?", "Say hello"]
            batches = list(runner._create_batches(prompts, batch_size=2))
            
            assert len(batches) == 2
            assert batches[0] == ["What is 1+1?", "Color of sky?"]
            assert batches[1] == ["Capital of USA?", "Say hello"]
    
    def test_create_batches_remainder(self, real_agent, simple_prompts):
        """Test batch creation with remainder"""
        with AgentRunner(real_agent) as runner:
            # Add one more prompt to create remainder
            prompts = simple_prompts + ["Extra prompt"]
            batches = list(runner._create_batches(prompts, batch_size=2))
            
            assert len(batches) == 3
            assert batches[0] == ["What is 1+1?", "Color of sky?"]
            assert batches[1] == ["Capital of USA?", "Say hello"]
            assert batches[2] == ["Extra prompt"]
    
    def test_create_batches_empty_list(self, real_agent):
        """Test batch creation with empty prompt list"""
        with AgentRunner(real_agent) as runner:
            batches = list(runner._create_batches([], batch_size=2))
            assert len(batches) == 0


class TestProcessBatch:
    """Test single batch processing"""
    
    @pytest.mark.asyncio
    async def test_process_batch_success(self, real_agent):
        """Test successful batch processing"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?", "Color of sky?"]
            results = await runner._process_batch(prompts)
            
            assert len(results) == 2
            # Check that we got actual results (not None)
            assert results[0] is not None
            assert results[1] is not None
            # Results should have an output attribute
            assert hasattr(results[0], 'output')
            assert hasattr(results[1], 'output')
    
    @pytest.mark.asyncio
    async def test_process_batch_with_framer(self, real_agent, simple_framer_func):
        """Test batch processing with framer function"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?", "Color of sky?"]
            result = await runner._process_batch(prompts, framer_func=simple_framer_func)
            
            assert isinstance(result, pd.DataFrame)
            assert len(result) == 2
            assert "prompt_index" in result.columns
            assert "response" in result.columns
            assert "success" in result.columns
            # Both should be successful
            assert all(result["success"])
    
    @pytest.mark.asyncio
    async def test_process_batch_with_output_type(self, real_agent):
        """Test batch processing with output type specification"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?"]
            results = await runner._process_batch(prompts, output_type=str)
            
            assert len(results) == 1
            assert results[0] is not None
            # Should still get a result with output
            assert hasattr(results[0], 'output')
    
    @pytest.mark.asyncio
    async def test_process_batch_with_timeout(self, real_agent):
        """Test batch processing respects timeout setting"""
        with AgentRunner(real_agent, timeout=0.001) as runner:  # Very short timeout
            prompts = ["What is 1+1?"]
            results = await runner._process_batch(prompts)
            
            # With such a short timeout, we expect the request to timeout and return None
            assert len(results) == 1
            assert results[0] is None  # Should timeout and return None


class TestDatabaseIntegration:
    """Test database integration functionality"""
    
    @pytest.mark.asyncio
    async def test_write_batch_to_db(self, real_agent, temp_db_file):
        """Test writing batch results to database"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Create test DataFrame
            df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
            semaphore = asyncio.Semaphore(1)
            
            # Should not raise an error
            await runner._write_batch_to_db(df, semaphore)
        
        # Verify data was written by reading it back
        import duckdb
        con = duckdb.connect(temp_db_file)
        try:
            result = con.execute("SELECT * FROM results").fetchall()
            assert len(result) == 2
            assert result[0] == (1, "a")
            assert result[1] == (2, "b")
        finally:
            con.close()
    
    @pytest.mark.asyncio
    async def test_write_batch_empty_dataframe(self, real_agent, temp_db_file):
        """Test writing empty DataFrame to database"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            empty_df = pd.DataFrame()
            semaphore = asyncio.Semaphore(1)
            
            # Should not raise an error and should return early
            await runner._write_batch_to_db(empty_df, semaphore)
        
        # Verify no table was created
        import duckdb
        con = duckdb.connect(temp_db_file)
        try:
            tables = con.execute("SHOW TABLES").fetchall()
            assert len(tables) == 0
        except:
            # If SHOW TABLES fails, that's also fine - no data was written
            pass
        finally:
            con.close()


class TestPydanticModelIntegration:
    """Test Pydantic model integration with and without framer functions"""
    
    @pytest.mark.asyncio
    async def test_pydantic_model_without_framer(self, real_agent, pydantic_response_model, pydantic_prompts):
        """Test using Pydantic model without framer function"""
        with AgentRunner(real_agent) as runner:
            # Use a single prompt to keep test simple and fast
            test_prompt = pydantic_prompts[0]
            results = await runner.run([test_prompt], output_type=pydantic_response_model)
        
        assert len(results) == 1
        assert results[0] is not None
        assert hasattr(results[0], 'output')
        
        # Verify the output is a Pydantic model instance
        output = results[0].output
        assert hasattr(output, 'model_dump')
        assert hasattr(output, 'answer')
        assert hasattr(output, 'thoughts')
        
        # Verify the structure
        model_dict = output.model_dump()
        assert 'answer' in model_dict
        assert 'thoughts' in model_dict
        assert isinstance(model_dict['answer'], str)
        assert isinstance(model_dict['thoughts'], list)
    
    @pytest.mark.asyncio
    async def test_pydantic_model_with_framer(self, real_agent, pydantic_response_model, pydantic_framer_func, pydantic_prompts):
        """Test using Pydantic model with framer function that uses model_dump()"""
        with AgentRunner(real_agent) as runner:
            # Use first two prompts to test batch processing
            test_prompts = pydantic_prompts[:2]
            result = await runner.run(
                test_prompts, 
                output_type=pydantic_response_model,
                framer_func=pydantic_framer_func,
                batch_size=2
            )
        
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 2
        
        # Verify DataFrame has expected columns from model_dump()
        expected_columns = {'prompt_index', 'success', 'answer', 'thoughts'}
        assert expected_columns.issubset(set(result.columns))
        
        # Verify all were successful
        assert all(result["success"])
        
        # Verify data types and content
        assert all(isinstance(answer, str) for answer in result["answer"])
        assert all(isinstance(thoughts, list) for thoughts in result["thoughts"])
        assert all(len(answer) > 0 for answer in result["answer"])
    
    @pytest.mark.asyncio
    async def test_pydantic_model_with_database_integration(self, real_agent, pydantic_response_model, pydantic_framer_func, pydantic_prompts, temp_db_file):
        """Test Pydantic model with database persistence using model_dump()"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Use single prompt for faster test
            test_prompt = pydantic_prompts[0]
            result = await runner.run(
                [test_prompt],
                output_type=pydantic_response_model,
                framer_func=pydantic_framer_func
            )
        
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 1
        
        # Verify database persistence
        import duckdb
        con = duckdb.connect(temp_db_file)
        try:
            stored_results = con.execute("SELECT * FROM results").fetchall()
            assert len(stored_results) == 1
            
            # Verify we have the structured data from Pydantic model (only model fields are stored)
            row = stored_results[0]
            # Should have exactly 2 columns: answer and thoughts (Pydantic model fields only)
            assert len(row) == 2
            
            # Verify we have the answer and thoughts (order may vary)
            answer_found = False
            thoughts_found = False
            for value in row:
                if isinstance(value, str) and len(value) > 0 and value != "FAILED" and '[' not in value:  # answer column
                    answer_found = True
                elif isinstance(value, str) and '[' in value:  # JSON-serialized thoughts
                    thoughts_found = True
            
            assert answer_found, "Answer column not found in database"
            assert thoughts_found, "Thoughts column not found in database"
            
        finally:
            con.close()


class TestRunMethod:
    """Test the main run method with various scenarios"""
    
    @pytest.mark.asyncio
    async def test_run_basic(self, real_agent, simple_prompts):
        """Test basic run functionality"""
        with AgentRunner(real_agent) as runner:
            results = await runner.run(simple_prompts, batch_size=2, max_concurrency=2)
        
        assert len(results) == len(simple_prompts)
        # All results should be valid (not None)
        assert all(result is not None for result in results)
        # All results should have output
        assert all(hasattr(result, 'output') for result in results)
    
    @pytest.mark.asyncio
    async def test_run_with_framer(self, real_agent, simple_prompts, simple_framer_func):
        """Test run with framer function"""
        with AgentRunner(real_agent) as runner:
            result = await runner.run(
                simple_prompts, 
                batch_size=2, 
                max_concurrency=2, 
                framer_func=simple_framer_func
            )
        
        assert isinstance(result, pd.DataFrame)
        assert len(result) == len(simple_prompts)
        assert "prompt_index" in result.columns
        assert "response" in result.columns
        assert "success" in result.columns
        # All should be successful
        assert all(result["success"])
    
    @pytest.mark.asyncio
    async def test_run_database_without_framer_error(self, real_agent, simple_prompts, temp_db_file):
        """Test that database mode requires framer function"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            with pytest.raises(ValueError, match="Database file .* is configured but framer_func is not provided"):
                await runner.run(simple_prompts)
    
    @pytest.mark.asyncio
    async def test_run_with_database(self, real_agent, simple_prompts, simple_framer_func, temp_db_file):
        """Test run with database persistence"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            result = await runner.run(
                simple_prompts,
                batch_size=2,
                max_concurrency=2,
                framer_func=simple_framer_func
            )
        
        assert isinstance(result, pd.DataFrame)
        assert len(result) == len(simple_prompts)
        
        # Verify data was persisted to database
        import duckdb
        con = duckdb.connect(temp_db_file)
        try:
            stored_results = con.execute("SELECT * FROM results").fetchall()
            assert len(stored_results) == len(simple_prompts)
        finally:
            con.close()
    
    @pytest.mark.asyncio
    async def test_run_empty_prompts(self, real_agent):
        """Test run with empty prompt list"""
        with AgentRunner(real_agent) as runner:
            results = await runner.run([])
        
        assert len(results) == 0
    
    @pytest.mark.asyncio
    async def test_run_single_prompt(self, real_agent):
        """Test run with single prompt"""
        with AgentRunner(real_agent) as runner:
            results = await runner.run(["What is 1+1?"])
        
        assert len(results) == 1
        assert results[0] is not None
        assert hasattr(results[0], 'output')
    
    @pytest.mark.asyncio
    async def test_run_without_context_manager_raises_error(self, real_agent, simple_prompts):
        """Test that run() raises RuntimeError when not used in context manager"""
        runner = AgentRunner(real_agent)
        
        with pytest.raises(RuntimeError, match="AgentRunner.run\\(\\) must be called within a context manager"):
            await runner.run(simple_prompts)


class TestErrorHandling:
    """Test error handling and edge cases"""
    
    @pytest.mark.asyncio
    async def test_run_with_invalid_prompts(self, real_agent):
        """Test handling when prompts might cause issues"""
        with AgentRunner(real_agent) as runner:
            # Test with various edge case prompts
            edge_prompts = [
                "",  # Empty string
                "What is 1+1?",  # Normal prompt
                "A" * 50,  # Long prompt (but not too long)
            ]
            
            results = await runner.run(edge_prompts, batch_size=2)
        
        # Should handle all prompts gracefully
        assert len(results) == len(edge_prompts)
        # At least the normal prompt should work
        assert any(result is not None for result in results)


class TestAdvancedErrorHandling:
    """Test advanced error handling scenarios and edge cases"""
    
    @pytest.mark.asyncio
    async def test_agent_run_failure_handling(self, real_agent, simple_framer_func):
        """Test handling when individual agent.run calls fail"""
        with AgentRunner(real_agent) as runner:
            # Create a mock agent that will fail on certain prompts
            import unittest.mock
            
            original_run = runner.agent.run
            
            async def failing_run(prompt, output_type=None):
                if "FAIL_THIS_PROMPT" in prompt:
                    raise Exception("Simulated agent failure")
                return await original_run(prompt, output_type=output_type)
            
            # Mock the agent's run method
            with unittest.mock.patch.object(runner.agent, 'run', side_effect=failing_run):
                prompts = ["What is 1+1?", "FAIL_THIS_PROMPT", "What is 2+2?"]
                
                # Test without framer (should return list with None for failed prompt)
                results = await runner.run(prompts, batch_size=2)
                
                assert len(results) == 3
                assert results[0] is not None  # First prompt should succeed
                assert results[1] is None      # Second prompt should fail (None)
                assert results[2] is not None  # Third prompt should succeed
                
                # Test with framer (should handle None results gracefully)
                result_df = await runner.run(prompts, batch_size=2, framer_func=simple_framer_func)
                
                assert isinstance(result_df, pd.DataFrame)
                assert len(result_df) == 3
                assert result_df.iloc[0]["success"] == True   # First should succeed
                assert result_df.iloc[1]["success"] == False  # Second should fail
                assert result_df.iloc[2]["success"] == True   # Third should succeed
                assert result_df.iloc[1]["response"] == "FAILED"  # Failed response
    
    @pytest.mark.asyncio 
    async def test_database_connection_failure(self, real_agent, simple_framer_func, temp_db_file):
        """Test handling when database connection fails"""
        with AgentRunner(real_agent, db_file="/invalid/path/database.db") as runner:
            # The AgentRunner handles database failures gracefully and returns empty DataFrame
            result = await runner.run(
                ["What is 1+1?"], 
                framer_func=simple_framer_func
            )
        
        # Should return empty DataFrame when all batches fail due to DB issues
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 0
    
    @pytest.mark.asyncio
    async def test_framer_function_failure(self, real_agent):
        """Test handling when framer function fails"""
        with AgentRunner(real_agent) as runner:
            def failing_framer(results):
                raise Exception("Framer function failed")
            
            # The AgentRunner handles framer failures and returns empty DataFrame
            result = await runner.run(
                ["What is 1+1?"], 
                framer_func=failing_framer
            )
        
        # Should return empty DataFrame when framer fails
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 0
    
    @pytest.mark.asyncio
    async def test_empty_dataframe_return_on_all_failures(self, real_agent, simple_framer_func):
        """Test that empty DataFrame is returned when all batches fail with framer"""
        with AgentRunner(real_agent) as runner:
            # Create a framer that always returns None (simulating all batch failures)
            def always_fail_framer(results):
                return simple_framer_func(results)
            
            # Mock the _process_batch method to always return None (simulating batch failures)
            import unittest.mock
            
            original_process_batch = runner._process_batch
            
            async def failing_process_batch(*args, **kwargs):
                # Simulate a scenario where _process_batch returns None due to failure
                return None
            
            with unittest.mock.patch.object(runner, '_process_batch', side_effect=failing_process_batch):
                result = await runner.run(
                    ["What is 1+1?", "What is 2+2?"], 
                    batch_size=1,
                    framer_func=simple_framer_func
                )
                
                # Should return empty DataFrame when all batches fail
                assert isinstance(result, pd.DataFrame)
                assert len(result) == 0
    
    @pytest.mark.asyncio
    async def test_database_write_with_corrupted_data(self, real_agent, temp_db_file):
        """Test database handling with corrupted/invalid DataFrame"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            def corrupted_framer(results):
                # Return a DataFrame with problematic data that might cause DB issues
                import pandas as pd
                import numpy as np
                
                return pd.DataFrame({
                    "col1": [float('inf'), float('-inf'), np.nan],  # Problematic float values
                    "col2": [None, None, None],  # All None values
                    "col3": ["test", "test2", "test3"]  # Normal strings
                })
            
            # This might fail during database write due to problematic data
            # The exact behavior depends on DuckDB's handling of inf/nan values
            try:
                result = await runner.run(
                    ["What is 1+1?"], 
                    framer_func=corrupted_framer
                )
                # If it succeeds, verify the result
                assert isinstance(result, pd.DataFrame)
            except Exception:
                # If it fails, that's also acceptable for this test
                # The important thing is that we don't crash the whole system
                pass
    
    @pytest.mark.asyncio
    async def test_database_transaction_rollback(self, real_agent, temp_db_file):
        """Test database rollback handling when transaction fails"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            def problematic_framer(results):
                # Create a DataFrame with a column name that will cause issues
                import pandas as pd
                return pd.DataFrame({
                    "normal_col": ["test"],
                    "SELECT * FROM": ["malicious"],  # Column name that might cause SQL issues
                })
            
            # Mock duckdb connection to simulate transaction failure
            import unittest.mock
            import duckdb
            
            def mock_connect(db_file):
                mock_con = unittest.mock.Mock()
                mock_con.begin = unittest.mock.Mock()
                mock_con.execute = unittest.mock.Mock(side_effect=Exception("SQL execution failed"))
                mock_con.rollback = unittest.mock.Mock()
                mock_con.commit = unittest.mock.Mock()
                mock_con.close = unittest.mock.Mock()
                return mock_con
            
            with unittest.mock.patch('duckdb.connect', side_effect=mock_connect):
                # This should trigger the rollback error handling path and return empty DataFrame
                result = await runner.run(
                    ["What is 1+1?"], 
                    framer_func=problematic_framer
                )
                
                # Should return empty DataFrame when database operations fail
                assert isinstance(result, pd.DataFrame)
                assert len(result) == 0
    
    @pytest.mark.asyncio
    async def test_database_rollback_also_fails(self, real_agent, temp_db_file):
        """Test handling when both database operation and rollback fail"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            def simple_framer(results):
                import pandas as pd
                return pd.DataFrame({"col": ["test"]})
            
            # Mock duckdb connection where both execute and rollback fail
            import unittest.mock
            import duckdb
            
            def mock_connect(db_file):
                mock_con = unittest.mock.Mock()
                mock_con.begin = unittest.mock.Mock()
                mock_con.execute = unittest.mock.Mock(side_effect=Exception("SQL execution failed"))
                mock_con.rollback = unittest.mock.Mock(side_effect=Exception("Rollback failed"))
                mock_con.commit = unittest.mock.Mock()
                mock_con.close = unittest.mock.Mock()
                return mock_con
            
            with unittest.mock.patch('duckdb.connect', side_effect=mock_connect):
                # This should handle the case where rollback also fails and return empty DataFrame
                result = await runner.run(
                    ["What is 1+1?"], 
                    framer_func=simple_framer
                )
                
                # Should return empty DataFrame when both operations fail
                assert isinstance(result, pd.DataFrame)
                assert len(result) == 0


class TestDatabaseSchemaValidation:
    """Test database schema validation and detailed error reporting"""
    
    @pytest.mark.asyncio
    async def test_database_schema_validation_error_with_detailed_info(self, real_agent, temp_db_file):
        """Test that DatabaseSchemaValidationError includes detailed schema and DataFrame info"""
        class TypeMismatchModel(BaseModel):
            id: int = Field(description="Integer ID")
            name: str = Field(description="String name")
            active: bool = Field(description="Boolean active status")
        
        def problematic_framer(results):
            """Create a DataFrame that will cause type conversion issues"""
            return pd.DataFrame({
                "id": ["not_an_integer", "also_not_integer"],  # String in integer column
                "name": ["valid_string", "another_string"],    # Valid strings
                "active": ["not_a_boolean", "also_not_bool"]   # String in boolean column
            })
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # This should trigger a DatabaseSchemaValidationError
            with pytest.raises(DatabaseSchemaValidationError) as exc_info:
                await runner.run(
                    ["Test prompt"],
                    output_type=TypeMismatchModel,
                    framer_func=problematic_framer
                )
            
            # Verify the error message contains detailed debugging information
            error_msg = str(exc_info.value)
            
            # Should contain the original error
            assert "Database type conversion failed" in error_msg
            
            # Should contain table schema information
            assert "Table schema:" in error_msg
            assert "id" in error_msg  # Column names should be present
            assert "INTEGER" in error_msg or "VARCHAR" in error_msg  # Data types should be present
            
            # Should contain DataFrame dtypes information
            assert "DataFrame dtypes:" in error_msg
            assert "object" in error_msg or "int64" in error_msg  # Pandas dtypes
            
            # Should contain DataFrame sample
            assert "DataFrame sample" in error_msg
            assert "not_an_integer" in error_msg or "valid_string" in error_msg  # Sample data
    
    @pytest.mark.asyncio
    async def test_database_schema_validation_with_pydantic_schema(self, real_agent, temp_db_file):
        """Test that Pydantic schema is properly used for table creation"""
        class ValidModel(BaseModel):
            count: int = Field(description="Integer count")
            message: str = Field(description="String message")
            tags: Optional[List[str]] = Field(description="List of string tags")
        
        def valid_framer(results):
            """Create a DataFrame that matches the Pydantic schema"""
            return pd.DataFrame({
                "count": [1, 2, 3],
                "message": ["test1", "test2", "test3"],
                "tags": [["tag1"], ["tag2", "tag3"], []]  # Lists should work with JSON type
            })
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # This should succeed without errors
            result = await runner.run(
                ["What is 1+1?", "What is 2+2?", "What is 3+3?"],
                output_type=ValidModel,
                framer_func=valid_framer,
                batch_size=3
            )
            
            assert isinstance(result, pd.DataFrame)
            assert len(result) == 3
            
            # Verify the data was written to database with correct schema
            import duckdb
            con = duckdb.connect(temp_db_file)
            try:
                # Check table schema
                schema_result = con.execute("DESCRIBE results").fetchall()
                schema_dict = {row[0]: row[1] for row in schema_result}
                
                # Verify Pydantic types were mapped correctly
                assert schema_dict["count"] == "INTEGER"
                assert schema_dict["message"] == "VARCHAR"
                assert schema_dict["tags"] == "JSON"  # Optional[List[str]] should map to JSON
                
                # Verify data was inserted correctly
                data = con.execute("SELECT * FROM results").fetchall()
                assert len(data) == 3
                
            finally:
                con.close()
    
    @pytest.mark.asyncio
    async def test_cancellation_on_schema_validation_error(self, real_agent, temp_db_file):
        """Test that DatabaseSchemaValidationError properly cancels all async jobs"""
        class ProblematicModel(BaseModel):
            value: int = Field(description="Integer value")
        
        def mixed_framer(results):
            """Create a DataFrame where some batches will fail type conversion"""
            # Use a deterministic approach - first batch succeeds, second fails
            if not hasattr(mixed_framer, 'call_count'):
                mixed_framer.call_count = 0
            mixed_framer.call_count += 1
            
            if mixed_framer.call_count == 2:  # Second call will fail
                return pd.DataFrame({
                    "value": ["not_an_integer"]  # This will cause conversion error
                })
            else:
                return pd.DataFrame({
                    "value": [42]  # This would succeed
                })
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Use multiple prompts to create multiple batches
            # At least one should trigger the validation error
            with pytest.raises(DatabaseSchemaValidationError):
                await runner.run(
                    ["Test 1", "Test 2", "Test 3", "Test 4", "Test 5"],
                    output_type=ProblematicModel,
                    framer_func=mixed_framer,
                    batch_size=1,  # Each prompt in its own batch
                    max_concurrency=3
                )
            
            # After the exception, all tasks should be cleaned up
            assert len(runner.active_tasks) == 0
    
    @pytest.mark.asyncio
    async def test_complex_pydantic_type_mappings(self, real_agent, temp_db_file):
        """Test edge cases in Pydantic type mapping"""
        from datetime import datetime
        from typing import Dict, Tuple, Any, Union
        from enum import Enum
        
        class Status(Enum):
            ACTIVE = "active"
            INACTIVE = "inactive"
        
        class NestedModel(BaseModel):
            name: str = Field(description="Nested model name")
            
        class ComplexModel(BaseModel):
            timestamp: datetime = Field(description="Datetime field")
            nested_dict: Dict[str, Any] = Field(description="Dict field")  
            tuple_data: Tuple[int, str] = Field(description="Tuple field")
            nested_model: NestedModel = Field(description="Nested Pydantic model")
            status_enum: Status = Field(description="Custom enum")
            union_field: Union[int, str] = Field(description="Union type")
            bytes_data: Optional[bytes] = Field(description="Bytes field")
        
        def complex_framer(results):
            """Create DataFrame with complex data types"""
            import json
            return pd.DataFrame({
                "timestamp": ["2024-01-01T12:00:00"],
                "nested_dict": [json.dumps({"key": "value"})],
                "tuple_data": [json.dumps([42, "test"])],
                "nested_model": [json.dumps({"name": "test_name"})],
                "status_enum": ["active"],
                "union_field": [json.dumps("string_value")],
                "bytes_data": [None]  # Optional field
            })
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            result = await runner.run(
                ["Test complex types"],
                output_type=ComplexModel,
                framer_func=complex_framer
            )
            
            assert isinstance(result, pd.DataFrame)
            assert len(result) == 1
            
            # Verify schema was created with proper type mappings
            import duckdb
            con = duckdb.connect(temp_db_file)
            try:
                schema_result = con.execute("DESCRIBE results").fetchall()
                schema_dict = {row[0]: row[1] for row in schema_result}
                
                # Verify complex type mappings
                assert schema_dict["timestamp"] == "TIMESTAMP"
                assert schema_dict["nested_dict"] == "JSON"  # Dict -> JSON
                assert schema_dict["tuple_data"] == "JSON"   # Tuple -> JSON
                assert schema_dict["nested_model"] == "JSON" # Pydantic model -> JSON
                assert schema_dict["status_enum"] == "VARCHAR" # Enum -> VARCHAR
                assert schema_dict["union_field"] == "JSON"  # Union -> JSON
                assert schema_dict["bytes_data"] == "BLOB"   # bytes -> BLOB
                
            finally:
                con.close()
    
    def test_map_python_type_to_duckdb_comprehensive(self, real_agent):
        """Test _map_python_type_to_duckdb method directly"""
        from typing import Dict, Tuple, Union
        from enum import Enum
        import datetime
        
        class CustomEnum(Enum):
            VALUE1 = "value1"
        
        class CustomPydanticModel(BaseModel):
            field: str
        
        with AgentRunner(real_agent) as runner:
            # Test all basic types
            assert runner._type_converter.map_python_type_to_db(int) == "INTEGER"
            assert runner._type_converter.map_python_type_to_db(float) == "DOUBLE"
            assert runner._type_converter.map_python_type_to_db(bool) == "BOOLEAN"
            assert runner._type_converter.map_python_type_to_db(str) == "VARCHAR"
            assert runner._type_converter.map_python_type_to_db(bytes) == "BLOB"
            assert runner._type_converter.map_python_type_to_db(type(None)) == "VARCHAR"
            
            # Test generic types
            assert runner._type_converter.map_python_type_to_db(List[str]) == "JSON"
            assert runner._type_converter.map_python_type_to_db(Dict[str, int]) == "JSON"
            assert runner._type_converter.map_python_type_to_db(Tuple[int, str]) == "JSON"
            assert runner._type_converter.map_python_type_to_db(Optional[str]) == "VARCHAR"  # Should unwrap Optional
            assert runner._type_converter.map_python_type_to_db(Union[int, str]) == "JSON"   # Complex union -> JSON
            
            # Test datetime types
            assert runner._type_converter.map_python_type_to_db(datetime.datetime) == "TIMESTAMP"
            assert runner._type_converter.map_python_type_to_db(datetime.date) == "TIMESTAMP"
            assert runner._type_converter.map_python_type_to_db(datetime.time) == "TIMESTAMP"
            
            # Test enum types
            assert runner._type_converter.map_python_type_to_db(CustomEnum) == "VARCHAR"
            
            # Test Pydantic model types
            assert runner._type_converter.map_python_type_to_db(CustomPydanticModel) == "JSON"
            
            # Test unknown types default to JSON
            class CustomClass:
                pass
            assert runner._type_converter.map_python_type_to_db(CustomClass) == "JSON"
    
    @pytest.mark.asyncio
    async def test_schema_creation_without_model_fields(self, real_agent, temp_db_file):
        """Test behavior when output_type doesn't have model_fields"""
        def simple_framer(results):
            return pd.DataFrame({"col": ["test"]})
        
        # Use a regular class instead of Pydantic model
        class NotAPydanticModel:
            pass
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Should fall back to SELECT-based table creation
            result = await runner.run(
                ["Test"],
                output_type=NotAPydanticModel,
                framer_func=simple_framer
            )
            assert isinstance(result, pd.DataFrame)
            assert len(result) == 1
            
            # Verify table was still created (via fallback method)
            import duckdb
            con = duckdb.connect(temp_db_file)
            try:
                data = con.execute("SELECT * FROM results").fetchall()
                assert len(data) == 1
            finally:
                con.close()
    
    @pytest.mark.asyncio
    async def test_schema_error_handling_when_describe_fails(self, real_agent, temp_db_file):
        """Test error handling when DESCRIBE table fails"""
        def problematic_framer(results):
            return pd.DataFrame({"id": ["not_an_integer"]})
        
        class SimpleModel(BaseModel):
            id: int = Field(description="Integer ID")
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Mock the connection to make DESCRIBE fail
            import unittest.mock
            import duckdb
            
            original_connect = duckdb.connect
            
            def mock_connect(db_file):
                mock_con = unittest.mock.Mock()
                mock_con.execute = unittest.mock.Mock()
                mock_con.close = unittest.mock.Mock()
                
                # Make DESCRIBE fail but other queries work initially
                def side_effect(query):
                    if "DESCRIBE" in query:
                        raise Exception("DESCRIBE failed")
                    elif "INSERT" in query:
                        raise Exception("Type conversion failed") 
                    elif "information_schema.tables" in query:
                        return unittest.mock.Mock(fetchone=lambda: [0])  # Table doesn't exist
                    else:
                        return unittest.mock.Mock(fetchall=lambda: [])
                
                mock_con.execute.side_effect = side_effect
                return mock_con
            
            with unittest.mock.patch('duckdb.connect', side_effect=mock_connect):
                with pytest.raises(DatabaseSchemaValidationError) as exc_info:
                    await runner.run(["Test"], output_type=SimpleModel, framer_func=problematic_framer)
                
                # Should still contain error info even when schema retrieval fails
                error_msg = str(exc_info.value)
                assert "Unable to retrieve table schema" in error_msg
                assert "DataFrame dtypes:" in error_msg  # Should still have DataFrame info
    
    @pytest.mark.asyncio
    async def test_concurrent_schema_validation_errors(self, real_agent, temp_db_file):
        """Test that schema validation errors from concurrent batches are handled properly"""
        class IntModel(BaseModel):
            value: int = Field(description="Integer value")
        
        call_count = 0
        def alternating_framer(results):
            nonlocal call_count
            call_count += 1
            if call_count % 2 == 0:
                # Every other batch will fail
                return pd.DataFrame({"value": ["not_an_integer"]})
            else:
                return pd.DataFrame({"value": [42]})
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            with pytest.raises(DatabaseSchemaValidationError):
                await runner.run(
                    ["Test 1", "Test 2", "Test 3", "Test 4"],
                    output_type=IntModel,
                    framer_func=alternating_framer,
                    batch_size=1,
                    max_concurrency=4  # High concurrency to test race conditions
                )
            
            # Verify cleanup occurred
            assert len(runner.active_tasks) == 0
    
    @pytest.mark.asyncio
    async def test_empty_dataframe_with_schema_validation(self, real_agent, temp_db_file):
        """Test schema validation with empty DataFrame"""
        def empty_framer(results):
            # Create DataFrame with correct columns but no data
            df = pd.DataFrame({"id": []})
            df["id"] = df["id"].astype('int64')  # Ensure correct dtype
            return df
        
        class SimpleModel(BaseModel):
            id: int = Field(description="Integer ID")
        
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            # Should handle empty DataFrame gracefully
            result = await runner.run(
                ["Test"],
                output_type=SimpleModel,
                framer_func=empty_framer
            )
            assert isinstance(result, pd.DataFrame)
            assert len(result) == 0
            
            # Verify table was created but has no data
            import duckdb
            con = duckdb.connect(temp_db_file)
            try:
                # Table should exist
                tables = con.execute("SHOW TABLES").fetchall()
                assert len(tables) > 0
                
                # But should have no data
                data = con.execute("SELECT * FROM results").fetchall()
                assert len(data) == 0
            finally:
                con.close()
    


class TestLogging:
    """Test logger functionality"""
    
    def test_logger_none_no_output(self, real_agent, capfd):
        """Test that no logging occurs when logger=None"""
        runner = AgentRunner(real_agent, logger=None)
        runner._log_info("This should not appear")
        
        captured = capfd.readouterr()
        assert captured.out == ""
        assert captured.err == ""
    
    def test_logger_with_mock_logger(self, real_agent):
        """Test that logging works when logger is provided"""
        from unittest.mock import Mock
        
        mock_logger = Mock()
        runner = AgentRunner(real_agent, logger=mock_logger)
        
        test_message = "Test log message"
        runner._log_info(test_message)
        
        mock_logger.info.assert_called_once_with(test_message)
    
    def test_logger_parameter_stored(self, real_agent):
        """Test that logger parameter is properly stored"""
        from unittest.mock import Mock
        
        mock_logger = Mock()
        runner = AgentRunner(real_agent, logger=mock_logger)
        
        assert runner.logger is mock_logger
        
        runner_no_logger = AgentRunner(real_agent, logger=None)
        assert runner_no_logger.logger is None

    @pytest.mark.asyncio
    async def test_batch_completion_includes_cost(self, real_agent):
        """Test that batch completion logging includes cost information"""
        from unittest.mock import Mock, patch
        
        mock_logger = Mock()
        
        def simple_framer(results):
            return pd.DataFrame({"response": [r.data if r else None for r in results]})
        
        # Mock the agent to return predictable usage stats
        with patch.object(real_agent, 'run') as mock_run, \
             patch.object(real_agent, 'get_usage') as mock_get_usage:
            
            # Mock result
            mock_result = Mock()
            mock_result.data = "test response"
            mock_run.return_value = mock_result
            
            # Mock usage stats
            mock_usage = pd.Series({
                'requests': 1,
                'request_tokens': 10,
                'response_tokens': 5,
                'thoughts_tokens': 0,
                'total_tokens': 15
            })
            mock_get_usage.return_value = mock_usage
            
            with AgentRunner(real_agent, logger=mock_logger) as runner:
                # Mock the get_usage method on the runner to return cost info
                with patch.object(runner, 'get_usage') as mock_runner_usage:
                    mock_runner_usage.return_value = pd.Series({
                        'requests': 1,
                        'request_tokens': 10,
                        'response_tokens': 5,
                        'thoughts_tokens': 0,
                        'total_tokens': 15,
                        'total_cost': 0.0250
                    })
                    
                    await runner.run(["test prompt"], framer_func=simple_framer)
            
            # Check that the logger was called with cost information
            log_calls = [call.args[0] for call in mock_logger.info.call_args_list]
            
            # Find the "completed batch:" message (new format)
            completed_batch_msgs = [msg for msg in log_calls if "completed batch:" in msg]
            assert len(completed_batch_msgs) > 0
            
            # Verify the cost is included in the message (new format uses 2 decimal places and static padding, no timestamp, left-justified)
            cost_msg = completed_batch_msgs[0]
            # With BATCH_WIDTH=6, batch numbers are now left-justified: "1      of 1     "
            assert "completed batch: 1      of 1     " in cost_msg
            assert "cost: $0.03   of $0.03  " in cost_msg
            # Verify no timestamp is present
            assert not any(char.isdigit() and cost_msg[i-1] == ':' and cost_msg[i+1] == ':' for i, char in enumerate(cost_msg[1:-1], 1))

    @pytest.mark.asyncio
    async def test_final_cost_summary_logging(self, real_agent):
        """Test that final cost summary is logged with breakdown by token type"""
        from unittest.mock import Mock, patch
        
        mock_logger = Mock()
        
        def simple_framer(results):
            return pd.DataFrame({"response": [r.data if r else None for r in results]})
        
        # Mock the agent to return predictable usage stats
        with patch.object(real_agent, 'run') as mock_run, \
             patch.object(real_agent, 'get_usage') as mock_get_usage:
            
            # Mock result
            mock_result = Mock()
            mock_result.data = "test response"
            mock_run.return_value = mock_result
            
            # Mock usage stats
            mock_usage = pd.Series({
                'requests': 2,
                'request_tokens': 20,
                'response_tokens': 10,
                'thoughts_tokens': 5,
                'total_tokens': 35
            })
            mock_get_usage.return_value = mock_usage
            
            with AgentRunner(real_agent, logger=mock_logger) as runner:
                # Mock the get_usage method on the runner to return cost info
                with patch.object(runner, 'get_usage') as mock_runner_usage:
                    mock_runner_usage.return_value = pd.Series({
                        'requests': 2,
                        'request_tokens': 20,
                        'response_tokens': 10,
                        'thoughts_tokens': 5,
                        'total_tokens': 35,
                        'input_cost': 0.0200,
                        'output_cost': 0.0100,
                        'thoughts_cost': 0.0050,
                        'total_cost': 0.0350
                    })
                    
                    await runner.run(["test prompt 1", "test prompt 2"], framer_func=simple_framer)
            
            # Check that the logger was called with cost summary
            log_calls = [call.args[0] for call in mock_logger.info.call_args_list]
            
            # Find the "Cost Summary" message
            cost_summary_msgs = [msg for msg in log_calls if "Cost Summary" in msg]
            assert len(cost_summary_msgs) == 1
            
            # Verify the cost breakdown is included in the message
            summary_msg = cost_summary_msgs[0]
            assert "Input: $0.0200" in summary_msg
            assert "Output: $0.0100" in summary_msg  
            assert "Thoughts: $0.0050" in summary_msg
            assert "Total: $0.0350" in summary_msg


class TestConcurrencyControl:
    """Test concurrency control mechanisms"""
    
    @pytest.mark.asyncio
    async def test_batching_behavior(self, real_agent, simple_prompts):
        """Test that batching works correctly with real agents"""
        with AgentRunner(real_agent) as runner:
            # Test with different batch sizes
            result1 = await runner.run(simple_prompts, batch_size=1, max_concurrency=2)
            result2 = await runner.run(simple_prompts, batch_size=4, max_concurrency=1)
        
        # Both should return the same number of results
        assert len(result1) == len(result2) == len(simple_prompts)
        # All results should be valid
        assert all(r is not None for r in result1)
        assert all(r is not None for r in result2)


@pytest.mark.integration 
class TestIntegration:
    """Integration tests that test the full workflow"""
    
    @pytest.mark.asyncio
    async def test_full_workflow_with_database(self, real_agent, simple_prompts, simple_framer_func, temp_db_file):
        """Test complete workflow with database persistence"""
        with AgentRunner(real_agent, db_file=temp_db_file) as runner:
            result = await runner.run(
                simple_prompts,
                batch_size=2,
                max_concurrency=2,
                framer_func=simple_framer_func
            )
        
        # Verify results
        assert isinstance(result, pd.DataFrame)
        assert len(result) == len(simple_prompts)
        assert "prompt_index" in result.columns
        assert "response" in result.columns
        assert "success" in result.columns
        
        # Verify database persistence
        import duckdb
        con = duckdb.connect(temp_db_file)
        try:
            stored_results = con.execute("SELECT * FROM results").fetchall()
            # Should have same number of results as prompts
            assert len(stored_results) == len(simple_prompts)
            
            # Verify all were successful (success is the 3rd column, index 2)
            assert all(row[2] == True for row in stored_results)  # success column
            
            # Verify we have response data (response is the 2nd column, index 1)
            assert all(row[1] and row[1] != "FAILED" for row in stored_results)  # response column
        finally:
            con.close()
    
    @pytest.mark.asyncio
    async def test_workflow_without_database(self, real_agent, simple_prompts, simple_framer_func):
        """Test complete workflow without database"""
        with AgentRunner(real_agent) as runner:
            result = await runner.run(
                simple_prompts,
                batch_size=2,
                max_concurrency=2,
                framer_func=simple_framer_func
            )
        
        # Verify results
        assert isinstance(result, pd.DataFrame)
        assert len(result) == len(simple_prompts)
        assert all(result["success"])


class TestUsageTracking:
    """Test usage tracking and aggregation functionality"""
    
    def test_usage_stats_initialization(self, real_agent):
        """Test that usage stats are properly initialized"""
        with AgentRunner(real_agent) as runner:
            usage = runner.get_usage()
            
            assert isinstance(usage, pd.Series)
            assert usage['requests'] == 0
            assert usage['request_tokens'] == 0
            assert usage['response_tokens'] == 0
            assert usage['thoughts_tokens'] == 0
            assert usage['total_tokens'] == 0
            assert usage['input_cost'] == 0.0
            assert usage['output_cost'] == 0.0
            assert usage['thoughts_cost'] == 0.0
            assert usage['total_cost'] == 0.0
    
    @pytest.mark.asyncio
    async def test_usage_aggregation_equals_sum_of_individual_calls(self, real_agent):
        """Test that aggregated usage equals the sum of individual agent.get_usage() calls"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?", "What is 2+2?", "What is 3+3?"]
            results = await runner.run(prompts, batch_size=1)
            
            # Verify we got all results
            assert len(results) == 3
            successful_results = [r for r in results if r is not None]
            
            # Calculate expected usage by summing individual agent.get_usage() calls
            expected_usage = None
            for result in successful_results:
                individual_usage = runner.agent.get_usage(result)
                if expected_usage is None:
                    expected_usage = individual_usage.copy()
                else:
                    expected_usage += individual_usage
            
            # Get actual aggregated usage from runner
            actual_usage = runner.get_usage()
            
            # Compare the core usage fields (excluding cost calculations)
            core_fields = ['requests', 'request_tokens', 'response_tokens', 'thoughts_tokens', 'total_tokens']
            for field in core_fields:
                assert actual_usage[field] == expected_usage[field], \
                    f"Mismatch in {field}: actual={actual_usage[field]}, expected={expected_usage[field]}"
    
    @pytest.mark.asyncio
    async def test_usage_aggregation_multiple_batches_equals_individual_sum(self, real_agent):
        """Test aggregation equals individual sum across multiple batches"""
        with AgentRunner(real_agent) as runner:
            prompts = ["What is 1+1?", "What is 2+2?", "What is 3+3?", "What is 4+4?"]
            results = await runner.run(prompts, batch_size=2, max_concurrency=2)
            
            # Get successful results
            successful_results = [r for r in results if r is not None]
            assert len(successful_results) > 0, "At least some results should succeed"
            
            # Calculate expected usage by manually summing individual calls
            expected_usage = None
            for result in successful_results:
                individual_usage = runner.agent.get_usage(result)
                if expected_usage is None:
                    expected_usage = individual_usage.copy()
                else:
                    expected_usage += individual_usage
            
            # Get actual aggregated usage
            actual_usage = runner.get_usage()
            
            # Verify aggregation matches manual sum
            core_fields = ['requests', 'request_tokens', 'response_tokens', 'thoughts_tokens', 'total_tokens']
            for field in core_fields:
                assert actual_usage[field] == expected_usage[field], \
                    f"Multi-batch mismatch in {field}: actual={actual_usage[field]}, expected={expected_usage[field]}"
    
    @pytest.mark.asyncio
    async def test_usage_partial_failures_aggregation_accuracy(self, real_agent):
        """Test that aggregation is accurate when some calls fail"""
        with AgentRunner(real_agent) as runner:
            # Create a mix of prompts, some which might timeout with very short timeout
            import unittest.mock
            
            # Mock agent.run to simulate some failures
            original_run = runner.agent.run
            call_count = 0
            
            async def selective_failing_run(prompt, output_type=None):
                nonlocal call_count
                call_count += 1
                # Fail every third call
                if call_count % 3 == 0:
                    raise Exception("Simulated failure")
                return await original_run(prompt, output_type=output_type)
            
            with unittest.mock.patch.object(runner.agent, 'run', side_effect=selective_failing_run):
                prompts = ["Test 1", "Test 2", "Test 3", "Test 4", "Test 5", "Test 6"]
                results = await runner.run(prompts, batch_size=2)
                
                # Get successful results
                successful_results = [r for r in results if r is not None]
                
                # Calculate expected usage from successful results only
                expected_usage = None
                for result in successful_results:
                    individual_usage = runner.agent.get_usage(result)
                    if expected_usage is None:
                        expected_usage = individual_usage.copy()
                    else:
                        expected_usage += individual_usage
                
                # Get actual aggregated usage
                actual_usage = runner.get_usage()
                
                # Verify aggregation matches (should only include successful calls)
                if expected_usage is not None:
                    core_fields = ['requests', 'request_tokens', 'response_tokens', 'thoughts_tokens', 'total_tokens']
                    for field in core_fields:
                        assert actual_usage[field] == expected_usage[field], \
                            f"Partial failure mismatch in {field}: actual={actual_usage[field]}, expected={expected_usage[field]}"
                else:
                    # All calls failed, usage should be zero
                    assert actual_usage['requests'] == 0
                    assert actual_usage['total_tokens'] == 0
    
    @pytest.mark.asyncio
    async def test_usage_aggregation_single_batch(self, real_agent):
        """Test usage aggregation for a single batch"""
        with AgentRunner(real_agent) as runner:
            # Run a single prompt
            results = await runner.run(["What is 1+1?"], batch_size=1)
            
            # Verify we got a result
            assert len(results) == 1
            assert results[0] is not None
            
            # Check usage stats were aggregated
            usage = runner.get_usage()
            assert usage['requests'] > 0
            assert usage['request_tokens'] > 0
            assert usage['response_tokens'] > 0
            assert usage['thoughts_tokens'] >= 0  # May be 0 if no thoughts used
            assert usage['total_tokens'] > 0
            
            # Verify the token relationship: total = request + response + thoughts
            assert usage['total_tokens'] == usage['request_tokens'] + usage['response_tokens'] + usage['thoughts_tokens']
    
    def test_usage_cost_calculations(self, real_agent):
        """Test cost calculations with auto-detected cost model"""
        # Test with default agent (should use cost model based on agent.model_name)
        with AgentRunner(real_agent) as runner:
            # Manually set some usage stats for testing
            import pandas as pd
            with runner._usage_lock:
                runner._usage_stats = pd.Series({
                    'requests': 5,
                    'request_tokens': 1000,
                    'response_tokens': 500,
                    'thoughts_tokens': 200,
                    'total_tokens': 1700
                })
            
            usage = runner.get_usage()
            # Should use the cost model that was auto-detected from the agent
            expected_input_cost = 1000 * runner._cost_model.input_ppm_cost / 1_000_000
            expected_output_cost = 500 * runner._cost_model.output_ppm_cost / 1_000_000
            expected_thoughts_cost = 200 * runner._cost_model.thought_ppm_cost / 1_000_000
            expected_total = expected_input_cost + expected_output_cost + expected_thoughts_cost
            
            assert abs(usage['input_cost'] - expected_input_cost) < 1e-10
            assert abs(usage['output_cost'] - expected_output_cost) < 1e-10
            assert abs(usage['thoughts_cost'] - expected_thoughts_cost) < 1e-10
            assert abs(usage['total_cost'] - expected_total) < 1e-10
    
    def test_usage_copy_independence(self, real_agent):
        """Test that get_usage() returns an independent copy"""
        with AgentRunner(real_agent) as runner:
            # Manually set some usage stats
            import pandas as pd
            with runner._usage_lock:
                runner._usage_stats = pd.Series({
                    'requests': 3,
                    'request_tokens': 100,
                    'response_tokens': 50,
                    'thoughts_tokens': 25,
                    'total_tokens': 175
                })
            
            # Get usage copy
            usage1 = runner.get_usage()
            usage2 = runner.get_usage()
            
            # Modify one copy
            usage1['requests'] = 999
            usage1['input_cost'] = 999.0
            usage1['thoughts_tokens'] = 999
            
            # Other copy should be unchanged
            assert usage2['requests'] == 3
            assert usage2['input_cost'] == 100 * runner._cost_model.input_ppm_cost / 1_000_000
            assert usage2['thoughts_tokens'] == 25
            
            # Original internal stats should be unchanged
            with runner._usage_lock:
                assert runner._usage_stats['requests'] == 3
                assert runner._usage_stats['thoughts_tokens'] == 25


@pytest.mark.integration 
class TestCleanupAndTimeoutFeatures:
    """Test new cleanup and timeout functionality"""
    
    @pytest.mark.asyncio
    async def test_manual_cleanup_functions(self, real_agent):
        """Test manual cleanup functions work correctly"""
        # Create runner but don't use context manager
        runner = AgentRunner(real_agent)
        
        # Simulate some running state
        runner._is_running = True
        
        # Test manual cleanup
        cancel_all_running_tasks()
        cleanup_all_agents()
        
        # Verify cleanup was called
        assert runner._cleanup_done is True
    
    @pytest.mark.asyncio
    async def test_timeout_configuration(self, real_agent):
        """Test timeout configuration is properly set"""
        with AgentRunner(real_agent, timeout=120.0) as runner:
            assert runner.timeout == 120.0
    
    @pytest.mark.asyncio
    async def test_task_tracking(self, real_agent, simple_prompts):
        """Test that tasks are properly tracked"""
        with AgentRunner(real_agent) as runner:
            # Initially no tasks
            assert len(runner.active_tasks) == 0
            
            # Start a run (but don't await it immediately)
            task = asyncio.create_task(runner.run(simple_prompts[:1], batch_size=1))
            
            # Give it a moment to start
            await asyncio.sleep(0.1)
            
            # Should have active tasks while running
            # (This test is somewhat timing-dependent)
            
            # Complete the task
            await task
    
    @pytest.mark.asyncio
    async def test_context_manager_cleanup_on_exception(self, real_agent):
        """Test context manager cleans up even when exception occurs"""
        runner = AgentRunner(real_agent)
        
        try:
            with runner:
                # Simulate some work
                runner._is_running = True
                # Raise an exception
                raise ValueError("Test exception")
        except ValueError:
            pass  # Expected
        
        # Cleanup should still have been called
        assert runner._cleanup_done is True
    
    @pytest.mark.asyncio
    async def test_timeout_prevents_hanging(self, real_agent):
        """Test that timeout prevents requests from hanging indefinitely"""
        # Use a very short timeout to test timeout behavior
        with AgentRunner(real_agent, timeout=0.001) as runner:
            # This should timeout quickly
            results = await runner.run(["Tell me a long story"], batch_size=1)
            
            # Should return None due to timeout
            assert len(results) == 1
            assert results[0] is None
    
    def test_global_task_tracker_singleton(self):
        """Test that TaskTracker is a proper singleton"""
        from easier.agent_runner import _task_tracker, TaskTracker
        
        # Create new instances
        tracker1 = TaskTracker()
        tracker2 = TaskTracker()
        
        # Should be the same instance
        assert tracker1 is tracker2
        assert tracker1 is _task_tracker
    
    @pytest.mark.asyncio
    async def test_multiple_runners_cleanup(self, real_agent):
        """Test cleanup works with multiple runner instances"""
        runners = []
        
        try:
            # Create multiple runners
            for i in range(3):
                runner = AgentRunner(real_agent, timeout=30.0)
                runners.append(runner)
                runner._is_running = True
            
            # Global cleanup should clean up all
            cleanup_all_agents()
            
            # All should be cleaned up
            for runner in runners:
                assert runner._cleanup_done is True
                
        finally:
            # Ensure cleanup
            for runner in runners:
                try:
                    runner._cleanup()
                except:
                    pass