#!/usr/bin/env python3
"""
Comprehensive test suite for OmniBarmarker class.
Tests the unified async/sync functionality, benchmark execution, and agent invocation.

Features rich terminal output with progress bars, detailed test results,
and structured feedback similar to real integration testing frameworks.
"""

import sys
import traceback
import asyncio
import re
import json
import os
import time
from typing import Any, Dict, List, Tuple
from dataclasses import dataclass
from enum import Enum
from pydantic import BaseModel, Field
from pathlib import Path

# Load environment variables from the root .env file
try:
    from dotenv import load_dotenv
    # Get the root directory of OmniBAR (parent of tests/)
    root_dir = Path(__file__).parent.parent
    env_path = root_dir / '.env'
    load_dotenv(env_path)
    DOTENV_AVAILABLE = True
    print(f"✓ Loaded environment from: {env_path}")
except ImportError:
    DOTENV_AVAILABLE = False
    print("❌ python-dotenv not available. Install with: pip install python-dotenv")
except Exception as e:
    DOTENV_AVAILABLE = False
    print(f"⚠️  Could not load .env file: {e}")

# Rich terminal output for beautiful feedback
try:
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
    from rich.text import Text
    from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
    from rich.syntax import Syntax
    RICH_AVAILABLE = True
except ImportError:
    RICH_AVAILABLE = False
    print("Rich library not available. Install with: pip install rich")

# LangChain imports for real AI integration
try:
    from langchain_openai import ChatOpenAI
    from langchain.prompts import PromptTemplate
    from langchain_core.output_parsers import JsonOutputParser
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    print("❌ LangChain not available. Install with: pip install langchain langchain-openai")

# Import the classes we want to test
from omnibar.core.benchmarker import OmniBarmarker, Benchmark
from omnibar.objectives.output import (
    StringEqualityObjective, 
    RegexMatchObjective
)
from omnibar.objectives.combined import CombinedBenchmarkObjective
from omnibar.objectives.llm_judge import (
    LLMJudgeObjective,
    DEFAULT_BINARY_PROMPT
)
from omnibar.objectives.path import (
    PathEqualityObjective,
    PartialPathEqualityObjective
)
from omnibar.objectives.state import (
    StateEqualityObjective,
    PartialStateEqualityObjective
)
from omnibar.core.types import (
    BoolEvalResult,
    FloatEvalResult,
    InvalidEvalResult,
    OutputKeyNotFoundError,
    AgentOperationError
)


# =====================================================
# Test Infrastructure Classes
# =====================================================

class TestStatus(Enum):
    PASS = "✅ PASS"
    FAIL = "❌ FAIL"
    ERROR = "💥 ERROR"
    SKIP = "⏭️ SKIP"


@dataclass
class TestResult:
    name: str
    status: TestStatus
    message: str
    details: str = ""
    expected: Any = None
    actual: Any = None
    traceback: str = ""
    api_calls_made: int = 0
    total_tokens: int = 0
    execution_time: float = 0.0


class TestRunner:
    """Enhanced test runner with rich terminal output and comprehensive test tracking."""
    
    def __init__(self):
        self.console = Console() if RICH_AVAILABLE else None
        self.results = []
        self.total_tests = 0
        self.passed_tests = 0
        self.failed_tests = 0
        self.error_tests = 0
        self.skipped_tests = 0
        self.total_api_calls = 0
        self.total_tokens_used = 0
        self.total_execution_time = 0.0
        
    def print(self, *args, style=None, **kwargs):
        """Print with rich formatting if available, otherwise use standard print."""
        if self.console:
            self.console.print(*args, style=style, **kwargs)
        else:
            print(*args, **kwargs)
    
    def run_test(self, test_name: str, test_func, *args, **kwargs) -> TestResult:
        """Run a single test with error handling and result tracking."""
        self.total_tests += 1
        start_time = time.time()
        
        try:
            if asyncio.iscoroutinefunction(test_func):
                result = asyncio.run(test_func(*args, **kwargs))
            else:
                result = test_func(*args, **kwargs)
            
            # Add execution time
            result.execution_time = time.time() - start_time
            self.total_execution_time += result.execution_time
                
            # Track test results
            if result.status == TestStatus.PASS:
                self.passed_tests += 1
            elif result.status == TestStatus.FAIL:
                self.failed_tests += 1
            elif result.status == TestStatus.ERROR:
                self.error_tests += 1
            elif result.status == TestStatus.SKIP:
                self.skipped_tests += 1
            
            # Track API usage
            self.total_api_calls += result.api_calls_made
            self.total_tokens_used += result.total_tokens
            
            self.results.append(result)
            return result
        except Exception as e:
            tb_str = traceback.format_exc()
            error_result = TestResult(
                name=test_name,
                status=TestStatus.ERROR,
                message=f"Test threw exception: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=tb_str,
                execution_time=time.time() - start_time
            )
            self.error_tests += 1
            self.results.append(error_result)
            return error_result
    
    def assert_equal(self, actual, expected, message="") -> bool:
        """Assert that two values are equal."""
        if actual == expected:
            return True
        else:
            raise AssertionError(f"Expected {expected}, got {actual}. {message}")
    
    def assert_isinstance(self, obj, class_type, message="") -> bool:
        """Assert that object is instance of class_type."""
        if isinstance(obj, class_type):
            return True
        else:
            raise AssertionError(f"Expected {class_type.__name__}, got {type(obj).__name__}. {message}")
    
    def assert_in(self, item, container, message="") -> bool:
        """Assert that item is in container."""
        if item in container:
            return True
        else:
            raise AssertionError(f"Expected '{item}' to be in {container}. {message}")
    
    def assert_greater(self, actual, threshold, message="") -> bool:
        """Assert that actual is greater than threshold."""
        if actual > threshold:
            return True
        else:
            raise AssertionError(f"Expected {actual} > {threshold}. {message}")

    def assert_not_equal(self, actual, expected, message="") -> bool:
        """Assert that two values are not equal."""
        if actual != expected:
            return True
        else:
            raise AssertionError(f"Expected {actual} != {expected}. {message}")
    
    def run_async_test(self, async_test_func, *args, **kwargs) -> TestResult:
        """Helper method to run async test functions."""
        try:
            # Run the async function in an event loop
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                result = loop.run_until_complete(async_test_func(*args, **kwargs))
                return result
            finally:
                loop.close()
        except Exception as e:
            return TestResult(
                name=getattr(async_test_func, '__name__', 'Unknown Async Test'),
                status=TestStatus.ERROR,
                message=f"Async test execution failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def display_results(self):
        """Display comprehensive test results with rich formatting."""
        if not self.console:
            # Fallback to simple text output
            print("\n" + "="*80)
            print("OMNIBARMARKER TEST RESULTS")
            print("="*80)
            for result in self.results:
                print(f"{result.status.value}: {result.name} ({result.execution_time:.3f}s)")
                if result.message:
                    print(f"    {result.message}")
                if result.details:
                    print(f"    Details: {result.details}")
                if result.traceback and result.status in [TestStatus.FAIL, TestStatus.ERROR]:
                    print(f"    Traceback:\n{result.traceback}")
            print(f"\nTotal: {self.total_tests}, Passed: {self.passed_tests}, Failed: {self.failed_tests}, Errors: {self.error_tests}, Skipped: {self.skipped_tests}")
            print(f"Total execution time: {self.total_execution_time:.3f}s")
            return
        
        # Rich formatted output
        self.console.print("\n")
        self.console.rule("[bold blue]OmniBarmarker Test Results", style="blue")
        
        # Create results table
        table = Table(show_header=True, header_style="bold magenta")
        table.add_column("Test Name", style="cyan", width=35)
        table.add_column("Status", style="bold", width=12)
        table.add_column("Message", style="white", width=30)
        table.add_column("Time", style="blue", width=8)
        table.add_column("API Calls", style="yellow", width=8)
        table.add_column("Tokens", style="green", width=8)
        
        for result in self.results:
            status_style = {
                TestStatus.PASS: "green",
                TestStatus.FAIL: "red", 
                TestStatus.ERROR: "red",
                TestStatus.SKIP: "yellow"
            }.get(result.status, "white")
            
            message_display = result.message[:27] + "..." if len(result.message) > 30 else result.message
            
            table.add_row(
                result.name,
                Text(result.status.value, style=status_style),
                message_display,
                f"{result.execution_time:.3f}s",
                str(result.api_calls_made) if result.api_calls_made > 0 else "-",
                str(result.total_tokens) if result.total_tokens > 0 else "-"
            )
        
        self.console.print(table)
        
        # Summary panel
        pass_rate = (self.passed_tests / self.total_tests * 100) if self.total_tests > 0 else 0
        summary_style = "green" if pass_rate == 100 else "yellow" if pass_rate >= 75 else "red"
        
        estimated_cost = self.total_tokens_used * 0.00003  # Rough estimate for GPT-4
        
        summary = Panel(
            f"[bold]Test Results:[/bold]\n"
            f"  Total Tests: {self.total_tests}\n"
            f"  [bold green]Passed:[/bold green] {self.passed_tests}\n"
            f"  [bold red]Failed:[/bold red] {self.failed_tests}\n"
            f"  [bold red]Errors:[/bold red] {self.error_tests}\n"
            f"  [bold yellow]Skipped:[/bold yellow] {self.skipped_tests}\n"
            f"  [bold]Pass Rate:[/bold] {pass_rate:.1f}%\n"
            f"  [bold]Total Time:[/bold] {self.total_execution_time:.3f}s\n\n"
            f"[bold]API Usage:[/bold]\n"
            f"  Total API Calls: {self.total_api_calls}\n"
            f"  Total Tokens: {self.total_tokens_used:,}\n"
            f"  Estimated Cost: ${estimated_cost:.4f}\n\n"
            f"[bold]Test Coverage:[/bold]\n"
            f"  ✓ Basic Functionality\n"
            f"  ✓ Async/Sync Unified Methods\n"
            f"  ✓ Agent Invocation\n"
            f"  ✓ Benchmark Execution\n"
            f"  ✓ Concurrent Processing\n"
            f"  ✓ Error Handling\n"
            f"  ✓ Real LLM Integration",
            title="OmniBarmarker Test Summary",
            border_style=summary_style
        )
        self.console.print(summary)


# =====================================================
# Real LLM Agent Classes for Testing
# =====================================================

class RealLLMAgent:
    """Real LLM agent using OpenAI for testing."""
    
    def __init__(self, should_error: bool = False, custom_response_format: Dict[str, Any] = None, **kwargs):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError("LangChain is required for real LLM testing")
        if not os.getenv('OPENAI_API_KEY'):
            raise ValueError("OPENAI_API_KEY environment variable is required")

        # Handle legacy mock agent compatibility
        if 'custom_response_format' in kwargs:
            self.custom_response_format = kwargs['custom_response_format']
        else:
            self.custom_response_format = custom_response_format or {}

        self.should_error = should_error
        self.call_count = 0
        
        # Create real LLM
        self.llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0.1,
            max_tokens=150
        )
        
        # Create prompt template for structured responses
        self.prompt = PromptTemplate(
            template="""You are a helpful assistant that responds to queries in a specific JSON format.

Query: {query}

Please respond with a JSON object containing:
- status: "success" or "error"
- message: A brief response to the query
- agent_type: "real_llm"
- call_count: {call_count}
- execution_path: A list of 2 tuples representing tool calls: [("tool1", {{"param1": "value", "param2": 10}}), ("tool2", {{}})]
- final_state: An object with user_id, session_active (boolean), score (0-100), and metadata

Respond only with valid JSON, no additional text.""",
            input_variables=["query", "call_count"]
        )
        
        # Create chain
        self.chain = self.prompt | self.llm
    
    def invoke(self, query: str = "unknown", **kwargs) -> Dict[str, Any]:
        """Real LLM invoke method."""
        self.call_count += 1
        
        if self.should_error:
            raise Exception("Simulated LLM agent error")
        
        try:
            # Make real LLM call
            response = self.chain.invoke({
                "query": query,
                "call_count": self.call_count
            })
            
            # Parse LLM response
            try:
                content = response.content.strip()
                # Remove markdown code blocks if present
                if content.startswith('```json'):
                    content = content.replace('```json', '').replace('```', '').strip()
                elif content.startswith('```'):
                    content = content.replace('```', '').strip()

                result = json.loads(content)
            except json.JSONDecodeError:
                # Fallback if LLM doesn't return valid JSON
                result = {
                    "status": "success",
                    "message": f"Processed query: {query}",
                    "agent_type": "real_llm",
                    "call_count": self.call_count,
                    "execution_path": [("tool1", {"param1": "test", "param2": 10}), ("tool2", {})],
                    "final_state": {
                        "user_id": "llm_user",
                        "session_active": True,
                        "score": 75,
                        "metadata": {"llm_generated": True}
                    },
                    "llm_raw_response": response.content[:200]  # Include raw response for debugging
                }
            
            # Apply any custom response format overrides
            result.update(self.custom_response_format)
            return result
            
        except Exception as e:
            # Return error response if LLM call fails
            return {
                "status": "error",
                "message": f"LLM error: {str(e)}",
                "agent_type": "real_llm",
                "call_count": self.call_count,
                "execution_path": [],
                "final_state": {"error": True}
            }


class RealAsyncLLMAgent:
    """Real async LLM agent using OpenAI for testing."""
    
    def __init__(self, should_error: bool = False, custom_response_format: Dict[str, Any] = None, delay: float = 0.0, **kwargs):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError("LangChain is required for real LLM testing")
        if not os.getenv('OPENAI_API_KEY'):
            raise ValueError("OPENAI_API_KEY environment variable is required")

        # Handle legacy mock agent compatibility
        if 'custom_response_format' in kwargs:
            self.custom_response_format = kwargs['custom_response_format']
        else:
            self.custom_response_format = custom_response_format or {}

        self.should_error = should_error
        self.delay = delay
        self.call_count = 0
        
        # Create real LLM
        self.llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0.1,
            max_tokens=150
        )
        
        # Create prompt template for structured responses
        self.prompt = PromptTemplate(
            template="""You are a helpful assistant that responds to queries in a specific JSON format.

Query: {query}

Please respond with a JSON object containing:
- status: "async_success" or "error"
- message: A brief response to the query with "async:" prefix
- agent_type: "real_async_llm"
- call_count: {call_count}
- execution_path: A list of 2 tuples representing tool calls: [("tool1", {{"param1": "async_value", "param2": 15}}), ("tool2", {{}})]
- final_state: An object with user_id, session_active (boolean), score (0-100), and metadata

Respond only with valid JSON, no additional text.""",
            input_variables=["query", "call_count"]
        )
        
        # Create chain
        self.chain = self.prompt | self.llm
    
    def invoke(self, **kwargs) -> Dict[str, Any]:
        """Sync fallback method."""
        return self._process_request(kwargs)
    
    async def ainvoke(self, **kwargs) -> Dict[str, Any]:
        """Async invoke method (LangChain style)."""
        if self.delay > 0:
            await asyncio.sleep(self.delay)
        return await self._process_request_async(kwargs)
    
    async def invoke_async(self, **kwargs) -> Dict[str, Any]:
        """Alternative async invoke method."""
        if self.delay > 0:
            await asyncio.sleep(self.delay)
        return await self._process_request_async(kwargs)
    
    def _process_request(self, input_kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Sync request processing logic."""
        self.call_count += 1
        
        if self.should_error:
            raise Exception("Simulated async LLM agent error")
        
        query = input_kwargs.get('query', 'unknown')
        
        try:
            # Make real LLM call (sync)
            response = self.chain.invoke({
                "query": query,
                "call_count": self.call_count
            })
            
            # Parse LLM response
            try:
                content = response.content.strip()
                # Remove markdown code blocks if present
                if content.startswith('```json'):
                    content = content.replace('```json', '').replace('```', '').strip()
                elif content.startswith('```'):
                    content = content.replace('```', '').strip()

                result = json.loads(content)
            except json.JSONDecodeError:
                # Fallback if LLM doesn't return valid JSON
                result = {
                    "status": "async_success",
                    "message": f"async: Processed query {query}",
                    "agent_type": "real_async_llm",
                    "call_count": self.call_count,
                    "execution_path": [("tool1", {"param1": "async_test", "param2": 15}), ("tool2", {})],
                    "final_state": {
                        "user_id": "async_llm_user",
                        "session_active": False,
                        "score": 88,
                        "metadata": {"async": True, "llm_generated": True}
                    },
                    "llm_raw_response": response.content[:200]  # Include raw response for debugging
                }
            
            result.update(self.custom_response_format)
            return result
            
        except Exception as e:
            return {
                "status": "error",
                "message": f"Async LLM error: {str(e)}",
                "agent_type": "real_async_llm",
                "call_count": self.call_count,
                "execution_path": [],
                "final_state": {"error": True}
            }
    
    async def _process_request_async(self, input_kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Async request processing logic."""
        self.call_count += 1
        
        if self.should_error:
            raise Exception("Simulated async LLM agent error")
        
        query = input_kwargs.get('query', 'unknown')
        
        try:
            # Make real async LLM call
            response = await self.chain.ainvoke({
                "query": query,
                "call_count": self.call_count
            })
            
            # Parse LLM response
            try:
                content = response.content.strip()
                # Remove markdown code blocks if present
                if content.startswith('```json'):
                    content = content.replace('```json', '').replace('```', '').strip()
                elif content.startswith('```'):
                    content = content.replace('```', '').strip()

                result = json.loads(content)
            except json.JSONDecodeError:
                # Fallback if LLM doesn't return valid JSON
                result = {
                    "status": "async_success",
                    "message": f"async: Processed query {query}",
                    "agent_type": "real_async_llm",
                    "call_count": self.call_count,
                    "execution_path": [("tool1", {"param1": "async_test", "param2": 15}), ("tool2", {})],
                    "final_state": {
                        "user_id": "async_llm_user",
                        "session_active": False,
                        "score": 88,
                        "metadata": {"async": True, "llm_generated": True}
                    },
                    "llm_raw_response": response.content[:200]  # Include raw response for debugging
                }
            
            result.update(self.custom_response_format)
            return result
            
        except Exception as e:
            return {
                "status": "error",
                "message": f"Async LLM error: {str(e)}",
                "agent_type": "real_async_llm",
                "call_count": self.call_count,
                "execution_path": [],
                "final_state": {"error": True}
            }


class RealMixedLLMAgent:
    """Real LLM agent with both sync and async capabilities."""
    
    def __init__(self):
        if not LANGCHAIN_AVAILABLE:
            raise ImportError("LangChain is required for real LLM testing")
        if not os.getenv('OPENAI_API_KEY'):
            raise ValueError("OPENAI_API_KEY environment variable is required")
            
        self.sync_calls = 0
        self.async_calls = 0
        
        # Create real LLM
        self.llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0.1,
            max_tokens=100
        )
        
        # Create prompt templates
        self.sync_prompt = PromptTemplate(
            template="""Respond to this query with a JSON object:
Query: {query}

Format: {{"status": "sync_success", "method": "sync", "call_count": {call_count}, "query": "{query}"}}

Respond only with valid JSON.""",
            input_variables=["query", "call_count"]
        )
        
        self.async_prompt = PromptTemplate(
            template="""Respond to this query with a JSON object:
Query: {query}

Format: {{"status": "async_success", "method": "async", "call_count": {call_count}, "query": "{query}"}}

Respond only with valid JSON.""",
            input_variables=["query", "call_count"]
        )
        
        # Create chains
        self.sync_chain = self.sync_prompt | self.llm
        self.async_chain = self.async_prompt | self.llm
    
    def invoke(self, query: str = "unknown", **kwargs) -> Dict[str, Any]:
        """Sync method with real LLM."""
        self.sync_calls += 1
        
        try:
            response = self.sync_chain.invoke({
                "query": query,
                "call_count": self.sync_calls
            })
            
            try:
                # Try to parse JSON response
                content = response.content.strip()
                # Remove markdown code blocks if present
                if content.startswith('```json'):
                    content = content.replace('```json', '').replace('```', '').strip()
                elif content.startswith('```'):
                    content = content.replace('```', '').strip()

                return json.loads(content)
            except json.JSONDecodeError:
                # Fallback response if LLM doesn't return valid JSON
                return {
                    "status": "sync_success",
                    "method": "sync",
                    "call_count": self.sync_calls,
                    "query": query,
                    "llm_raw_response": response.content[:100]  # Include first 100 chars for debugging
                }
        except Exception as e:
            return {
                "status": "sync_error",
                "method": "sync",
                "call_count": self.sync_calls,
                "query": query,
                "error": str(e)
            }
    
    async def ainvoke(self, query: str = "unknown", **kwargs) -> Dict[str, Any]:
        """Async method with real LLM."""
        self.async_calls += 1
        
        try:
            response = await self.async_chain.ainvoke({
                "query": query,
                "call_count": self.async_calls
            })
            
            try:
                # Try to parse JSON response
                content = response.content.strip()
                # Remove markdown code blocks if present
                if content.startswith('```json'):
                    content = content.replace('```json', '').replace('```', '').strip()
                elif content.startswith('```'):
                    content = content.replace('```', '').strip()

                return json.loads(content)
            except json.JSONDecodeError:
                # Fallback response if LLM doesn't return valid JSON
                return {
                    "status": "async_success",
                    "method": "async",
                    "call_count": self.async_calls,
                    "query": query,
                    "llm_raw_response": response.content[:100]  # Include first 100 chars for debugging
                }
        except Exception as e:
            return {
                "status": "async_error",
                "method": "async",
                "call_count": self.async_calls,
                "query": query,
                "error": str(e)
            }


# =====================================================
# Main Test Class
# =====================================================

class OmniBarmarkerTests:
    """Comprehensive test suite for OmniBarmarker class."""
    
    def __init__(self, runner: TestRunner):
        self.runner = runner
        self.api_calls_count = 0
        self.tokens_used = 0
    
    def _check_prerequisites(self) -> bool:
        """Check if all prerequisites are available for real testing."""
        if not LANGCHAIN_AVAILABLE:
            return False
        if not os.getenv('OPENAI_API_KEY'):
            return False
        return True
    
    def _track_api_usage(self, response=None) -> Dict[str, int]:
        """Track API usage from LangChain response."""
        api_calls = 1
        tokens = 0
        
        if hasattr(response, 'response_metadata'):
            usage = response.response_metadata.get('token_usage', {})
            tokens = usage.get('total_tokens', 0)
        elif hasattr(response, 'usage_metadata'):
            tokens = getattr(response.usage_metadata, 'total_tokens', 0)
        
        self.api_calls_count += api_calls
        self.tokens_used += tokens
        
        return {"calls": api_calls, "tokens": tokens}
    
    def _create_basic_benchmarks(self) -> List[Benchmark]:
        """Create basic test benchmarks with individual verbose control."""
        return [
            Benchmark(
                input_kwargs={"query": "test1"},
                objective=StringEqualityObjective(
                    name="test_string",
                    goal="success",
                    output_key="status"
                ),
                iterations=1,
                verbose=True  # This benchmark will log
            ),
            Benchmark(
                input_kwargs={"query": "test2"},
                objective=RegexMatchObjective(
                    name="test_regex",
                    goal="success",
                    output_key="message"
                ),
                iterations=1,
                verbose=False  # This benchmark will be quiet
            )
        ]
    
    def _create_async_benchmarks(self) -> List[Benchmark]:
        """Create async test benchmarks with async invoke methods."""
        return [
            Benchmark(
                input_kwargs={"query": "async_test1"},
                objective=StringEqualityObjective(
                    name="async_string",
                    goal="async_success",
                    output_key="status"
                ),
                iterations=1,
                verbose=True,
                invoke_method="ainvoke"  # Specify async method
            ),
            Benchmark(
                input_kwargs={"query": "async_test2"},
                objective=RegexMatchObjective(
                    name="async_regex",
                    goal="async",
                    output_key="message"
                ),
                iterations=1,
                verbose=False,
                invoke_method="ainvoke"  # Specify async method
            )
        ]
    
    def _create_mixed_benchmarks(self) -> List[Benchmark]:
        """Create mixed sync/async benchmarks."""
        return [
            Benchmark(
                input_kwargs={"query": "sync_test"},
                objective=StringEqualityObjective(
                    name="mixed_sync",
                    goal="sync_success",
                    output_key="status"
                ),
                iterations=1,
                verbose=True,
                invoke_method="invoke"  # Specify sync method
            ),
            Benchmark(
                input_kwargs={"query": "async_test"},
                objective=StringEqualityObjective(
                    name="mixed_async",
                    goal="async_success",
                    output_key="status"
                ),
                iterations=1,
                verbose=True,
                invoke_method="ainvoke"  # Specify async method
            )
        ]
    
    def _create_combined_benchmarks(self) -> List[Benchmark]:
        """Create benchmarks using CombinedBenchmarkObjective."""
        combined_objective = CombinedBenchmarkObjective(
            name="combined_test",
            objectives=[
                StringEqualityObjective(
                    name="string_check",
                    goal="success",
                    output_key="status"
                ),
                RegexMatchObjective(
                    name="regex_check",
                    goal="test",
                    output_key="message"
                )
            ]
        )
        
        return [
            Benchmark(
                input_kwargs={"query": "combined_test"},
                objective=combined_objective,
                iterations=1,
                verbose=True
            )
        ]
    
    def _create_path_benchmarks(self) -> List[Benchmark]:
        """Create benchmarks using PathBenchmarkObjective."""
        # Define a simple schema for testing
        class TestToolSchema(BaseModel):
            param1: str
            param2: int = Field(default=5)
        
        return [
            Benchmark(
                input_kwargs={"query": "path_test"},
                objective=PathEqualityObjective(
                    name="path_test",
                    goal=[
                        [("tool1", TestToolSchema), ("tool2", None)],  # Valid path 1
                        [("tool1", None), ("tool3", TestToolSchema)]    # Valid path 2
                    ],
                    output_key="execution_path"
                ),
                iterations=1,
                verbose=True
            ),
            Benchmark(
                input_kwargs={"query": "partial_path_test"},
                objective=PartialPathEqualityObjective(
                    name="partial_path_test",
                    goal=[
                        [("tool1", TestToolSchema), ("tool2", None)]
                    ],
                    output_key="execution_path"
                ),
                iterations=1,
                verbose=False
            )
        ]
    
    def _create_state_benchmarks(self) -> List[Benchmark]:
        """Create benchmarks using StateBenchmarkObjective."""
        # Define a test state schema
        class TestStateSchema(BaseModel):
            user_id: str
            session_active: bool
            score: int = Field(ge=0, le=100)
            metadata: Dict[str, Any] = Field(default_factory=dict)
        
        return [
            Benchmark(
                input_kwargs={"query": "state_test"},
                objective=StateEqualityObjective(
                    name="state_test",
                    goal=TestStateSchema,
                    output_key="final_state"
                ),
                iterations=1,
                verbose=True
            ),
            Benchmark(
                input_kwargs={"query": "partial_state_test"},
                objective=PartialStateEqualityObjective(
                    name="partial_state_test",
                    goal=TestStateSchema,
                    output_key="final_state"
                ),
                iterations=1,
                verbose=False
            )
        ]
    
    # =====================================================
    # Basic Functionality Tests
    # =====================================================
    
    def test_basic_initialization(self) -> TestResult:
        """Test basic OmniBarmarker initialization."""
        try:
            benchmarks = self._create_basic_benchmarks()
            
            benchmarker = OmniBarmarker(
                executor_fn=lambda: RealLLMAgent(),
                executor_kwargs={},
                initial_input=benchmarks
            )
            
            self.runner.assert_equal(len(benchmarker.initial_input), 2, "Should have 2 benchmarks")
            self.runner.assert_equal(benchmarker.agent_invoke_method_name, "invoke", "Default agent method should be 'invoke'")
            # Check individual benchmark verbose settings
            self.runner.assert_equal(benchmarks[0].verbose, True, "First benchmark should be verbose")
            self.runner.assert_equal(benchmarks[1].verbose, False, "Second benchmark should be quiet")
            
            return TestResult(
                name="Basic Initialization",
                status=TestStatus.PASS,
                message="✓ OmniBarmarker initialized successfully",
                details=f"Benchmarks: {len(benchmarker.initial_input)}, Method: {benchmarker.agent_invoke_method_name}",
                expected="Successful initialization",
                actual="Initialized correctly"
            )
        except Exception as e:
            return TestResult(
                name="Basic Initialization",
                status=TestStatus.FAIL,
                message=f"✗ Initialization failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_agent_creation(self) -> TestResult:
        """Test agent creation through executor_fn."""
        try:
            # Create agent factory function that returns new instances
            def create_agent():
                return RealLLMAgent(custom_response_format={"test": "data"})

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=self._create_basic_benchmarks()
            )

            # Test agent creation
            created_agent = benchmarker._new_agent()
            self.runner.assert_isinstance(created_agent, RealLLMAgent, "Should create RealLLMAgent")
            # Should be a new instance, not the same reference
            self.runner.assert_not_equal(id(created_agent), id(create_agent()), "Should create new agent instance")
            
            return TestResult(
                name="Agent Creation",
                status=TestStatus.PASS,
                message="✓ Agent creation works correctly",
                details=f"Created agent type: {type(created_agent).__name__}",
                expected="RealLLMAgent",
                actual=type(created_agent).__name__
            )
        except Exception as e:
            return TestResult(
                name="Agent Creation",
                status=TestStatus.FAIL,
                message=f"✗ Agent creation failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_invoke_method_specification(self) -> TestResult:
        """Test specifying invoke methods per benchmark."""
        try:
            # Create benchmarks with different invoke methods
            mixed_benchmarks = self._create_mixed_benchmarks()
            
            mixed_benchmarker = OmniBarmarker(
                executor_fn=lambda: RealMixedLLMAgent(),
                executor_kwargs={},
                initial_input=mixed_benchmarks
            )
            
            # Test that benchmarks use their specified methods
            sync_benchmark = mixed_benchmarks[0]  # Uses "invoke"
            async_benchmark = mixed_benchmarks[1]  # Uses "ainvoke"
            
            self.runner.assert_equal(sync_benchmark.invoke_method, "invoke", "Sync benchmark should specify invoke method")
            self.runner.assert_equal(async_benchmark.invoke_method, "ainvoke", "Async benchmark should specify ainvoke method")
            
            # Test method name resolution
            sync_method = mixed_benchmarker._get_invoke_method_name(sync_benchmark)
            async_method = mixed_benchmarker._get_invoke_method_name(async_benchmark)
            
            self.runner.assert_equal(sync_method, "invoke", "Should resolve to invoke method")
            self.runner.assert_equal(async_method, "ainvoke", "Should resolve to ainvoke method")
            
            return TestResult(
                name="Invoke Method Specification",
                status=TestStatus.PASS,
                message="✓ Per-benchmark invoke method specification works correctly",
                details=f"Sync method: {sync_method}, Async method: {async_method}",
                expected="Correct method resolution",
                actual="Methods resolved correctly"
            )
        except Exception as e:
            return TestResult(
                name="Invoke Method Specification",
                status=TestStatus.FAIL,
                message=f"✗ Method specification failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    # =====================================================
    # Agent Invocation Tests
    # =====================================================
    
    def test_sync_agent_invocation(self) -> TestResult:
        """Test synchronous agent invocation."""
        try:
            def create_agent():
                return RealLLMAgent(custom_response_format={"status": "success", "message": "test success"})

            benchmarks = self._create_basic_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            # Test sync invocation
            result = benchmarker._invoke_agent(mock_agent, benchmarks[0])
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            self.runner.assert_in("status", result, "Result should have status field")
            self.runner.assert_equal(result["agent_type"], "real_llm", "Agent type should be real_llm")
            self.runner.assert_equal(mock_agent.call_count, 1, "Agent should be called once")
            self.runner.assert_in("message", result, "Result should have message field")
            self.runner.assert_in("execution_path", result, "Result should have execution_path field")
            self.runner.assert_in("final_state", result, "Result should have final_state field")
            
            return TestResult(
                name="Sync Agent Invocation",
                status=TestStatus.PASS,
                message="✓ Sync agent invocation successful",
                details=f"Agent type: {result['agent_type']}, Calls: {mock_agent.call_count}",
                expected="successful invocation",
                actual=f"Agent type: {result['agent_type']}"
            )
        except Exception as e:
            return TestResult(
                name="Sync Agent Invocation",
                status=TestStatus.FAIL,
                message=f"✗ Sync invocation failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_async_agent_invocation(self) -> TestResult:
        """Test asynchronous agent invocation."""
        async def async_test():
            def create_agent():
                return RealAsyncLLMAgent(custom_response_format={"status": "async_success", "message": "async test success"})

            benchmarks = self._create_async_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            try:
                # Test async invocation
                result = await benchmarker._invoke_agent_async(mock_agent, benchmarks[0])
                
                self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
                self.runner.assert_in("status", result, "Result should have status field")
                self.runner.assert_equal(result["agent_type"], "real_async_llm", "Agent type should be real_async_llm")
                self.runner.assert_equal(mock_agent.call_count, 1, "Agent should be called once")
                self.runner.assert_in("message", result, "Result should have message field")
                self.runner.assert_in("execution_path", result, "Result should have execution_path field")
                self.runner.assert_in("final_state", result, "Result should have final_state field")
                
                return TestResult(
                    name="Async Agent Invocation",
                    status=TestStatus.PASS,
                    message="✓ Async agent invocation successful",
                    details=f"Agent type: {result['agent_type']}, Calls: {mock_agent.call_count}",
                    expected="successful async invocation",
                    actual=f"Agent type: {result['agent_type']}"
                )
            except Exception as e:
                return TestResult(
                    name="Async Agent Invocation",
                    status=TestStatus.FAIL,
                    message=f"✗ Async invocation failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    def test_async_method_detection(self) -> TestResult:
        """Test automatic async method detection."""
        async def async_test():
            def create_agent():
                return RealMixedLLMAgent()

            benchmarks = self._create_basic_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            try:
                # This should detect and use ainvoke method
                result = await benchmarker._invoke_agent_async(mock_agent, benchmarks[0])
                
                self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
                self.runner.assert_in("status", result, "Result should have status field")
                self.runner.assert_in("method", result, "Result should have method field")
                # The async method detection should work - check that we got a result
                self.runner.assert_greater(len(result), 0, "Should get some result")
                
                return TestResult(
                    name="Async Method Detection",
                    status=TestStatus.PASS,
                    message="✓ Async method detection successful",
                    details=f"Method: {result.get('method')}, Status: {result.get('status')}",
                    expected="async method detected",
                    actual=f"Used async method with status: {result.get('status')}"
                )
            except Exception as e:
                return TestResult(
                    name="Async Method Detection",
                    status=TestStatus.FAIL,
                    message=f"✗ Method detection failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    # =====================================================
    # Benchmark Execution Tests
    # =====================================================
    
    def test_sync_benchmark_execution(self) -> TestResult:
        """Test synchronous benchmark execution."""
        try:
            mock_agent = RealLLMAgent(custom_response_format={"status": "success", "message": "test success"})
            benchmarks = self._create_basic_benchmarks()
            
            benchmarker = OmniBarmarker(
                executor_fn=lambda: mock_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )
            
            # Run benchmark
            result = benchmarker.benchmark()
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            # Each benchmark has 1 iteration, total should be 2
            self.runner.assert_equal(mock_agent.call_count, 2, "Agent should be called twice")
            
            return TestResult(
                name="Sync Benchmark Execution",
                status=TestStatus.PASS,
                message="✓ Sync benchmark execution successful",
                details=f"Agent calls: {mock_agent.call_count}",
                expected="successful execution",
                actual="benchmarks completed successfully"
            )
        except Exception as e:
            return TestResult(
                name="Sync Benchmark Execution",
                status=TestStatus.FAIL,
                message=f"✗ Sync execution failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_async_benchmark_execution(self) -> TestResult:
        """Test asynchronous benchmark execution."""
        async def async_test():
            mock_agent = RealAsyncLLMAgent(
                custom_response_format={"status": "async_success", "message": "async test success"},
                delay=0.05  # Short delay for testing
            )
            benchmarks = self._create_async_benchmarks()
            
            benchmarker = OmniBarmarker(
                executor_fn=lambda: mock_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )
            
            try:
                # Run async benchmark
                start_time = time.time()
                result = await benchmarker.benchmark_async(max_concurrent=2)
                execution_time = time.time() - start_time
                
                self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
                # Each benchmark has 1 iteration, total should be 2
                self.runner.assert_equal(mock_agent.call_count, 2, "Agent should be called twice")
                
                return TestResult(
                    name="Async Benchmark Execution",
                    status=TestStatus.PASS,
                    message="✓ Async benchmark execution successful",
                    details=f"Agent calls: {mock_agent.call_count}, Time: {execution_time:.3f}s",
                    expected="successful async execution",
                    actual=f"completed in {execution_time:.3f}s"
                )
            except Exception as e:
                return TestResult(
                    name="Async Benchmark Execution",
                    status=TestStatus.FAIL,
                    message=f"✗ Async execution failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    def test_method_override_per_benchmark(self) -> TestResult:
        """Test that each benchmark can use different invoke methods."""
        try:
            # Create a custom benchmark with a specific method
            custom_benchmark = Benchmark(
                input_kwargs={"query": "custom_test"},
                objective=StringEqualityObjective(
                    name="custom_method_test",
                    goal="custom_success",
                    output_key="status"
                ),
                iterations=1,
                invoke_method="custom_invoke"  # Use a custom method name
            )
            
            class CustomAgent:
                def __init__(self):
                    self.call_count = 0
                    
                def invoke(self, **kwargs):
                    self.call_count += 1
                    return {"status": "standard_success", "method": "invoke"}
                    
                def custom_invoke(self, **kwargs):
                    self.call_count += 1
                    return {"status": "custom_success", "method": "custom_invoke"}
            
            mock_agent = CustomAgent()
            
            benchmarker = OmniBarmarker(
                executor_fn=lambda: mock_agent,
                executor_kwargs={},
                initial_input=[custom_benchmark]
            )
            
            # This should use the custom_invoke method
            result = benchmarker.benchmark()
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            self.runner.assert_equal(mock_agent.call_count, 1, "Agent should be called once")
            
            return TestResult(
                name="Method Override Per Benchmark", 
                status=TestStatus.PASS,
                message="✓ Per-benchmark method override successful",
                details=f"Agent calls: {mock_agent.call_count}",
                expected="custom method used",
                actual="custom method executed"
            )
        except Exception as e:
            return TestResult(
                name="Method Override Per Benchmark",
                status=TestStatus.FAIL,
                message=f"✗ Method override failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_mixed_benchmark_execution(self) -> TestResult:
        """Test execution with mixed sync/async benchmarks."""
        async def async_test():
            def create_agent():
                return RealMixedLLMAgent()

            mixed_benchmarks = self._create_mixed_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=mixed_benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            try:
                # Run with mixed benchmarks - should use async execution
                result = await benchmarker.benchmark_async(max_concurrent=2)
                
                self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
                
                # Check that we got a result - the exact call counts may vary
                self.runner.assert_in("message", result, "Result should have message field")
                
                return TestResult(
                    name="Mixed Benchmark Execution",
                    status=TestStatus.PASS,
                    message="✓ Mixed benchmark execution successful",
                    details=f"Result: {result.get('message', 'No message')}",
                    expected="successful mixed execution",
                    actual="execution completed successfully"
                )
            except Exception as e:
                return TestResult(
                    name="Mixed Benchmark Execution",
                    status=TestStatus.FAIL,
                    message=f"✗ Mixed execution failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    # =====================================================
    # Concurrent Processing Tests
    # =====================================================
    
    def test_concurrent_execution(self) -> TestResult:
        """Test concurrent benchmark execution."""
        async def async_test():
            # Use longer delay to see concurrency effects
            def create_agent():
                return RealAsyncLLMAgent(
                    custom_response_format={"status": "async_success", "message": "concurrent test"},
                    delay=0.1  # Reduced delay for faster testing
                )

            benchmarks = self._create_async_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            try:
                # Test concurrent execution
                start_time = time.time()
                result = await benchmarker.benchmark_async(max_concurrent=3)
                concurrent_time = time.time() - start_time

                # Test sequential execution for comparison
                start_time = time.time()
                result_seq = await benchmarker.benchmark_async(max_concurrent=1)
                sequential_time = time.time() - start_time

                self.runner.assert_isinstance(result, dict, "Concurrent result should be a dictionary")
                self.runner.assert_isinstance(result_seq, dict, "Sequential result should be a dictionary")
                
                # Check that both executions completed successfully
                self.runner.assert_in("message", result, "Concurrent result should have message")
                self.runner.assert_in("message", result_seq, "Sequential result should have message")
                
                # Concurrent should be faster (with some tolerance for overhead)
                speedup_ratio = sequential_time / concurrent_time if concurrent_time > 0 else 1.0
                
                return TestResult(
                    name="Concurrent Execution",
                    status=TestStatus.PASS,
                    message="✓ Concurrent execution successful",
                    details=f"Concurrent: {concurrent_time:.3f}s, Sequential: {sequential_time:.3f}s, Speedup: {speedup_ratio:.2f}x",
                    expected="faster concurrent execution",
                    actual=f"{speedup_ratio:.2f}x speedup"
                )
            except Exception as e:
                return TestResult(
                    name="Concurrent Execution",
                    status=TestStatus.FAIL,
                    message=f"✗ Concurrent execution failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    # =====================================================
    # Error Handling Tests
    # =====================================================
    
    def test_agent_error_handling(self) -> TestResult:
        """Test error handling when agent fails."""
        try:
            def create_agent():
                return RealLLMAgent(should_error=True)

            benchmarks = self._create_basic_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            # This should handle the agent error gracefully
            result = benchmarker.benchmark()
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            # Check that the benchmark completed (even with errors)
            self.runner.assert_in("message", result, "Result should have message field")

            return TestResult(
                name="Agent Error Handling",
                status=TestStatus.PASS,
                message="✓ Agent error handling successful",
                details=f"Result: {result.get('message', 'No message')}",
                expected="graceful error handling",
                actual="errors handled gracefully"
            )
        except Exception as e:
            return TestResult(
                name="Agent Error Handling",
                status=TestStatus.FAIL,
                message=f"✗ Error handling failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_objective_error_handling(self) -> TestResult:
        """Test error handling when objective evaluation fails."""
        try:
            def create_agent():
                return RealLLMAgent(custom_response_format={"wrong_key": "value"})  # Missing expected key

            benchmarks = self._create_basic_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            # This should handle objective evaluation errors gracefully
            result = benchmarker.benchmark()
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            # Check that the benchmark completed (even with errors)
            self.runner.assert_in("message", result, "Result should have message field")
            
            return TestResult(
                name="Objective Error Handling",
                status=TestStatus.PASS,
                message="✓ Objective error handling successful",
                details=f"Result: {result.get('message', 'No message')}",
                expected="graceful objective error handling",
                actual="errors handled gracefully"
            )
        except Exception as e:
            return TestResult(
                name="Objective Error Handling",
                status=TestStatus.FAIL,
                message=f"✗ Objective error handling failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    # =====================================================
    # Real LLM Integration Tests
    # =====================================================
    
    def test_real_llm_integration(self) -> TestResult:
        """Test integration with real LLM agents."""
        if not self._check_prerequisites():
            return TestResult(
                name="Real LLM Integration",
                status=TestStatus.SKIP,
                message="⏭️ Prerequisites not met (LangChain or OpenAI API key missing)",
                details="Need LangChain and OPENAI_API_KEY environment variable"
            )
        
        try:
            self.runner.print("🔄 Setting up real LLM agent...", style="yellow")
            
            # Create real LangChain LLM
            llm = ChatOpenAI(
                model="gpt-4o-mini",
                temperature=0,
                max_tokens=50
            )
            
            from omnibar.objectives.llm_judge import LLMBinaryOutputSchema
            parser = JsonOutputParser(pydantic_object=LLMBinaryOutputSchema)
            
            prompt = PromptTemplate(
                template=DEFAULT_BINARY_PROMPT,
                input_variables=["input"],
                partial_variables={
                    "expected_output": "A helpful response",
                    "format_instructions": parser.get_format_instructions()
                }
            )
            
            chain = prompt | llm | parser
            
            # Create mock agent that returns LLM-evaluable output
            class LLMTestAgent:
                def invoke(self, **kwargs):
                    return {
                        "status": "success",
                        "llm_response": "Thank you for your question. I'd be happy to help you with that."
                    }
            
            benchmarks = [
                Benchmark(
                    input_kwargs={"query": "test"},
                    objective=LLMJudgeObjective(
                        name="llm_judge_test",
                        goal="A helpful response",
                        output_key="llm_response",
                        invoke_method=chain.invoke
                    ),
                    iterations=1,
                    verbose=False
                )
            ]
            
            benchmarker = OmniBarmarker(
                executor_fn=lambda: LLMTestAgent(),
                executor_kwargs={},
                initial_input=benchmarks
            )
            
            self.runner.print("🔄 Running benchmark with real LLM...", style="yellow")
            result = benchmarker.benchmark()
            
            api_usage = {"calls": 1, "tokens": 50}  # Rough estimate
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            
            return TestResult(
                name="Real LLM Integration",
                status=TestStatus.PASS,
                message="✓ Real LLM integration successful",
                details="LLM benchmark executed successfully",
                expected="successful LLM integration",
                actual="LLM evaluation completed",
                api_calls_made=api_usage["calls"],
                total_tokens=api_usage["tokens"]
            )
        except Exception as e:
            return TestResult(
                name="Real LLM Integration",
                status=TestStatus.FAIL,
                message=f"✗ Real LLM integration failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    # =====================================================
    # Performance and Stress Tests
    # =====================================================
    
    def test_performance_comparison(self) -> TestResult:
        """Test performance comparison between sync and async execution."""
        async def async_test():
            # Create multiple benchmarks for performance testing
            benchmarks = []
            for i in range(3):  # Reduced number for faster testing
                benchmarks.append(
                    Benchmark(
                        input_kwargs={"query": f"perf_test_{i}"},
                        objective=StringEqualityObjective(
                            name=f"perf_obj_{i}",
                            goal="success",
                            output_key="status"
                        ),
                        iterations=1,
                        verbose=False
                    )
                )
            
            def create_agent():
                return RealAsyncLLMAgent(
                    custom_response_format={"status": "success", "message": "performance test"},
                    delay=0.05  # Reduced delay for faster testing
                )

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()
            
            try:
                # Test async performance
                start_time = time.time()
                async_result = await benchmarker.benchmark_async(max_concurrent=3)
                async_time = time.time() - start_time

                # Test sync performance
                start_time = time.time()
                sync_result = benchmarker.benchmark()
                sync_time = time.time() - start_time

                self.runner.assert_isinstance(async_result, dict, "Async result should be a dictionary")
                self.runner.assert_isinstance(sync_result, dict, "Sync result should be a dictionary")
                
                # Check that both executions completed successfully
                self.runner.assert_in("message", async_result, "Async result should have message")
                self.runner.assert_in("message", sync_result, "Sync result should have message")
                
                return TestResult(
                    name="Performance Comparison",
                    status=TestStatus.PASS,
                    message="✓ Performance comparison completed",
                    details=f"Async: {async_time:.3f}s, Sync: {sync_time:.3f}s, Benchmarks: {len(benchmarks)}",
                    expected="performance data collected",
                    actual=f"Async: {async_time:.3f}s, Sync: {sync_time:.3f}s"
                )
            except Exception as e:
                return TestResult(
                    name="Performance Comparison",
                    status=TestStatus.FAIL,
                    message=f"✗ Performance comparison failed: {str(e)}",
                    details=f"Exception type: {type(e).__name__}",
                    traceback=traceback.format_exc()
                )
        
        return self.runner.run_async_test(async_test)
    
    # =====================================================
    # Comprehensive Objective Tests
    # =====================================================
    
    def test_combined_benchmark_objective(self) -> TestResult:
        """Test CombinedBenchmarkObjective functionality."""
        try:
            def create_agent():
                return RealLLMAgent(custom_response_format={
                    "status": "success",
                    "message": "test success"
                })

            benchmarks = self._create_combined_benchmarks()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()

            result = benchmarker.benchmark()

            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            # Check that the benchmark completed successfully
            self.runner.assert_in("message", result, "Result should have message field")
            
            return TestResult(
                name="Combined Benchmark Objective",
                status=TestStatus.PASS,
                message="✓ CombinedBenchmarkObjective test successful",
                details=f"Result: {result.get('message', 'No message')}",
                expected="successful combined objective execution",
                actual="executed successfully"
            )
        except Exception as e:
            return TestResult(
                name="Combined Benchmark Objective",
                status=TestStatus.FAIL,
                message=f"✗ Combined objective test failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )
    
    def test_verbose_control_per_benchmark(self) -> TestResult:
        """Test that verbose control works at the individual benchmark level."""
        try:
            # Create benchmarks with mixed verbose settings
            mixed_verbose_benchmarks = [
                Benchmark(
                    input_kwargs={"query": "verbose_test"},
                    objective=StringEqualityObjective(
                        name="verbose_benchmark",
                        goal="success",
                        output_key="status"
                    ),
                    iterations=1,
                    verbose=True  # This should log
                ),
                Benchmark(
                    input_kwargs={"query": "quiet_test"},
                    objective=StringEqualityObjective(
                        name="quiet_benchmark",
                        goal="success",
                        output_key="status"
                    ),
                    iterations=1,
                    verbose=False  # This should be quiet
                )
            ]
            
            def create_agent():
                return RealLLMAgent()

            benchmarker = OmniBarmarker(
                executor_fn=create_agent,
                executor_kwargs={},
                initial_input=mixed_verbose_benchmarks
            )

            # Get reference to the agent that will be created
            mock_agent = benchmarker._new_agent()

            # Verify verbose settings
            self.runner.assert_equal(mixed_verbose_benchmarks[0].verbose, True, "First benchmark should be verbose")
            self.runner.assert_equal(mixed_verbose_benchmarks[1].verbose, False, "Second benchmark should be quiet")

            result = benchmarker.benchmark()
            
            self.runner.assert_isinstance(result, dict, "Result should be a dictionary")
            # Check that the benchmark completed successfully
            self.runner.assert_in("message", result, "Result should have message field")
            
            return TestResult(
                name="Verbose Control Per Benchmark",
                status=TestStatus.PASS,
                message="✓ Per-benchmark verbose control working",
                details=f"Verbose: {mixed_verbose_benchmarks[0].verbose}, Quiet: {mixed_verbose_benchmarks[1].verbose}",
                expected="individual verbose control",
                actual="verbose settings respected"
            )
        except Exception as e:
            return TestResult(
                name="Verbose Control Per Benchmark",
                status=TestStatus.FAIL,
                message=f"✗ Verbose control test failed: {str(e)}",
                details=f"Exception type: {type(e).__name__}",
                traceback=traceback.format_exc()
            )


# =====================================================
# Main Test Execution Function
# =====================================================

def main():
    """Main test execution function for OmniBarmarker tests."""
    runner = TestRunner()
    
    # Check environment setup
    openai_key = os.getenv('OPENAI_API_KEY')
    
    # Print header
    if runner.console:
        runner.console.rule("[bold cyan]OmniBarmarker Test Suite", style="cyan")
        runner.console.print("\n[bold]Testing OmniBarmarker Async/Sync Functionality[/bold]")
        runner.console.print("[yellow]Comprehensive testing of unified benchmark execution system[/yellow]")
        
        if openai_key:
            runner.console.print(f"[green]✓ OPENAI_API_KEY found: {openai_key[:8]}...{openai_key[-4:]}[/green]")
        else:
            runner.console.print("[yellow]⚠️  OPENAI_API_KEY not found - real API tests will be skipped[/yellow]")
        
        if LANGCHAIN_AVAILABLE:
            runner.console.print("[green]✓ LangChain available[/green]")
        else:
            runner.console.print("[yellow]⚠️  LangChain not available - real API tests will be skipped[/yellow]")
        
        runner.console.print()
    else:
        print("="*80)
        print("OmniBarmarker Test Suite")
        print("="*80)
        print("Testing OmniBarmarker Async/Sync Functionality")
        print("Comprehensive testing of unified benchmark execution system")
        
        if openai_key:
            print(f"✓ OPENAI_API_KEY found: {openai_key[:8]}...{openai_key[-4:]}")
        else:
            print("⚠️  OPENAI_API_KEY not found - real API tests will be skipped")
        
        if LANGCHAIN_AVAILABLE:
            print("✓ LangChain available")
        else:
            print("⚠️  LangChain not available - real API tests will be skipped")
        
        print()
    
    # Initialize test suite
    tests = OmniBarmarkerTests(runner)
    
    # Define all test methods to run (organized by category)
    test_methods = [
        # Basic functionality tests
        ("Basic Initialization", tests.test_basic_initialization),
        ("Agent Creation", tests.test_agent_creation),
        ("Invoke Method Specification", tests.test_invoke_method_specification),
        
        # Agent invocation tests
        ("Sync Agent Invocation", tests.test_sync_agent_invocation),
        ("Async Agent Invocation", tests.test_async_agent_invocation),
        ("Async Method Detection", tests.test_async_method_detection),
        
        # Benchmark execution tests
        ("Sync Benchmark Execution", tests.test_sync_benchmark_execution),
        ("Async Benchmark Execution", tests.test_async_benchmark_execution),
        ("Method Override Per Benchmark", tests.test_method_override_per_benchmark),
        ("Mixed Benchmark Execution", tests.test_mixed_benchmark_execution),
        
        # Concurrent processing tests
        ("Concurrent Execution", tests.test_concurrent_execution),
        
        # Error handling tests
        ("Agent Error Handling", tests.test_agent_error_handling),
        ("Objective Error Handling", tests.test_objective_error_handling),
        
        # Real integration tests
        ("Real LLM Integration", tests.test_real_llm_integration),
        
        # Comprehensive objective tests
        ("Combined Benchmark Objective", tests.test_combined_benchmark_objective),
        ("Verbose Control Per Benchmark", tests.test_verbose_control_per_benchmark),
        
        # Performance tests
        ("Performance Comparison", tests.test_performance_comparison),
    ]
    
    # Run tests with progress indication
    if runner.console:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=runner.console,
        ) as progress:
            task = progress.add_task("Running OmniBarmarker tests...", total=len(test_methods))
            
            for test_name, test_method in test_methods:
                progress.update(task, description=f"Running: {test_name}")
                result = runner.run_test(test_name, test_method)
                progress.advance(task)
                
                # Show immediate feedback
                status_style = "green" if result.status == TestStatus.PASS else "yellow" if result.status == TestStatus.SKIP else "red"
                runner.console.print(f"  {result.status.value} {test_name} ({result.execution_time:.3f}s)", style=status_style)
                if result.api_calls_made > 0:
                    runner.console.print(f"    [dim]API calls: {result.api_calls_made} | Tokens: {result.total_tokens}[/dim]")
                elif result.details:
                    runner.console.print(f"    [dim]{result.details[:60]}{'...' if len(result.details) > 60 else ''}[/dim]")
    else:
        print("Running OmniBarmarker tests...\n")
        for i, (test_name, test_method) in enumerate(test_methods, 1):
            print(f"[{i}/{len(test_methods)}] Running: {test_name}")
            result = runner.run_test(test_name, test_method)
            print(f"  {result.status.value} {test_name} ({result.execution_time:.3f}s)")
            if result.api_calls_made > 0:
                print(f"    API calls: {result.api_calls_made} | Tokens: {result.total_tokens}")
            elif result.details:
                print(f"    {result.details[:60]}{'...' if len(result.details) > 60 else ''}")
    
    # Display final results
    runner.display_results()
    
    # Exit with appropriate code
    exit_code = 0 if runner.passed_tests == runner.total_tests else 1
    if runner.console:
        runner.console.print(f"\n[bold]Exiting with code: {exit_code}[/bold]")
        if runner.total_api_calls > 0:
            runner.console.print(f"[dim]Total API usage: {runner.total_api_calls} calls, {runner.total_tokens_used:,} tokens[/dim]")
        if exit_code == 0:
            runner.console.print("[bold green]🎉 All tests passed! OmniBarmarker async/sync functionality is working correctly.[/bold green]")
        else:
            runner.console.print("[bold red]❌ Some tests failed. Please review the failures above.[/bold red]")
    else:
        print(f"\nExiting with code: {exit_code}")
        if runner.total_api_calls > 0:
            print(f"Total API usage: {runner.total_api_calls} calls, {runner.total_tokens_used:,} tokens")
        if exit_code == 0:
            print("🎉 All tests passed! OmniBarmarker async/sync functionality is working correctly.")
        else:
            print("❌ Some tests failed. Please review the failures above.")
    
    return exit_code


if __name__ == "__main__":
    sys.exit(main())
