import asyncio
import time
from collections.abc import AsyncGenerator, Generator

import pytest
from timbal import Agent, Tool
from timbal.types.events import Event, OutputEvent
from timbal.types.message import Message

# ==============================================================================
# Test Handler Functions
# ==============================================================================

def sync_handler(x: str) -> str:
    """Simple synchronous handler for testing."""
    return f"sync:{x}"


def sync_gen_handler(count: int) -> Generator[int, None, None]:
    """Synchronous generator handler for testing."""
    for i in range(count):
        yield i


async def async_handler(x: str) -> str:
    """Simple asynchronous handler for testing."""
    await asyncio.sleep(0.01)  # Small delay to test async behavior
    return f"async:{x}"


async def async_gen_handler(count: int) -> AsyncGenerator[int, None]:
    """Asynchronous generator handler for testing."""
    for i in range(count):
        await asyncio.sleep(0.01)
        yield i


def error_handler(message: str) -> str:
    """Handler that raises an exception for testing error handling."""
    raise ValueError(f"Test error: {message}")


async def async_error_handler(message: str) -> str:
    """Async handler that raises an exception for testing error handling."""
    await asyncio.sleep(0.01)
    raise ValueError(f"Test async error: {message}")


# ==============================================================================
# Test Tools
# ==============================================================================

@pytest.fixture
def simple_tool():
    """Create a simple tool for testing."""
    return Tool(name="simple", handler=sync_handler)


@pytest.fixture
def async_tool():
    """Create an async tool for testing."""
    return Tool(name="async", handler=async_handler)


@pytest.fixture
def gen_tool():
    """Create a generator tool for testing."""
    return Tool(name="gen", handler=sync_gen_handler)


@pytest.fixture
def async_gen_tool():
    """Create an async generator tool for testing."""
    return Tool(name="async_gen", handler=async_gen_handler)


@pytest.fixture
def error_tool():
    """Create a tool that raises errors for testing."""
    return Tool(name="error", handler=error_handler)


# ==============================================================================
# Test Agents
# ==============================================================================

@pytest.fixture
def simple_agent():
    """Create a simple agent for testing."""
    def get_time() -> str:
        return time.strftime("%Y-%m-%d %H:%M:%S")
    
    return Agent(
        name="simple_agent",
        model="openai/gpt-4o-mini",
        tools=[get_time]
    )


@pytest.fixture
def multi_tool_agent():
    """Create an agent with multiple tools for testing."""
    def add(a: int, b: int) -> int:
        return a + b
    
    def multiply(a: int, b: int) -> int:
        return a * b
    
    def greet(name: str) -> str:
        return f"Hello, {name}!"
    
    return Agent(
        name="multi_tool_agent",
        model="openai/gpt-4o-mini",
        tools=[add, multiply, greet]
    )


@pytest.fixture
def math_agent():
    """Create an agent specialized for math operations."""
    def calculate(expression: str) -> str:
        """Safely evaluate simple math expressions."""
        try:
            # Only allow basic math operations for safety
            allowed_chars = set('0123456789+-*/().= ')
            if all(c in allowed_chars for c in expression):
                result = eval(expression)
                return f"{expression} = {result}"
            else:
                return "Invalid expression - only basic math allowed"
        except Exception as e:
            return f"Error: {str(e)}"
    
    return Agent(
        name="math_agent",
        model="openai/gpt-4o-mini",
        tools=[calculate],
        system_prompt="You are a helpful math assistant. Use the calculate tool for any math operations."
    )




# ==============================================================================
# Assertion Helpers
# ==============================================================================

def assert_event_sequence(events: list[Event], expected_types: list[type]):
    """Assert that events follow the expected type sequence."""
    assert len(events) == len(expected_types), f"Expected {len(expected_types)} events, got {len(events)}"
    for event, expected_type in zip(events, expected_types, strict=False):
        assert isinstance(event, expected_type), f"Expected {expected_type.__name__}, got {type(event).__name__}"


def assert_has_output_event(output: OutputEvent):
    """Assert that we have a valid OutputEvent."""
    assert isinstance(output, OutputEvent), f"Expected OutputEvent, got {type(output)}"


def assert_no_errors(output: OutputEvent):
    """Assert that the output contains no errors."""
    if output.error:
        pytest.fail(f"Found error in OutputEvent: {output.error}")


def skip_if_agent_error(output: OutputEvent, test_name: str = ""):
    """Skip test if agent execution failed - indicates implementation issue."""
    if output.error is not None:
        error_msg = output.error.get('message', str(output.error))
        pytest.skip(f"Agent execution failed in {test_name} - needs implementation fix: {error_msg}")


def assert_message_content(message: Message, expected_content: str = None):
    """Assert message properties and optionally check content."""
    assert isinstance(message, Message), f"Expected Message, got {type(message)}"
    assert message.role in ["user", "assistant", "tool"], f"Invalid role: {message.role}"
    if expected_content:
        content_str = str(message.content)
        assert expected_content.lower() in content_str.lower(), f"Expected '{expected_content}' in message content"


# ==============================================================================
# Performance Testing Utilities
# ==============================================================================

class Timer:
    """Simple timer context manager for performance testing."""
    
    def __init__(self):
        self.start_time = None
        self.end_time = None
        self.elapsed = None
    
    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        self.elapsed = self.end_time - self.start_time


# ==============================================================================
# Parametrized Test Data
# ==============================================================================

# Tool handler types for parametrized tests
TOOL_HANDLERS = [
    ("sync", sync_handler),
    ("async", async_handler),
    ("sync_gen", sync_gen_handler),
    ("async_gen", async_gen_handler),
]

# Error scenarios for testing
ERROR_SCENARIOS = [
    ("sync_error", error_handler),
    ("async_error", async_error_handler),
]

# Agent models for testing
TEST_MODELS = [
    "openai/gpt-4o-mini",
    "openai/gpt-4",
    "anthropic/claude-3-sonnet",
]

# Common test prompts
TEST_PROMPTS = [
    Message.validate({"role": "user", "content": "Hello, how are you?"}),
    Message.validate({"role": "user", "content": "What is 2 + 2?"}),
    Message.validate({"role": "user", "content": "Tell me a short joke."}),
]