"""Testing utilities for Logos ADK agents."""
from __future__ import annotations

import contextlib
from typing import Any

from logos_adk.agent import Agent
from logos_adk.errors import GuardrailError
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, StructuredResponse


def mock_agent(
    name: str,
    responses: list[str | Response] | None = None,
    stream_responses: list[str | list[str] | Response] | None = None,
) -> tuple[Agent, MockProvider]:
    """Create an Agent with a MockProvider for testing.

    Returns:
        Tuple of (agent, mock_provider).
    """
    provider = MockProvider(responses=responses, stream_responses=stream_responses)
    agent = Agent(name, model=provider)
    return agent, provider


def assert_response(response: Response, **checks: Any) -> None:
    """Assert response matches the given checks."""
    content = response.content

    if "contains" in checks:
        for word in checks["contains"]:
            if word not in content:
                raise AssertionError(
                    f"Response missing '{word}'. Content: {content[:200]}"
                )

    if "not_contains" in checks:
        for word in checks["not_contains"]:
            if word in content:
                raise AssertionError(
                    f"Response should not contain '{word}'. Content: {content[:200]}"
                )

    if "max_length" in checks:
        if len(content) > checks["max_length"]:
            raise AssertionError(
                f"Response length {len(content)} exceeds max_length {checks['max_length']}"
            )

    if "min_length" in checks:
        if len(content) < checks["min_length"]:
            raise AssertionError(
                f"Response length {len(content)} below min_length {checks['min_length']}"
            )

    if "schema" in checks:
        schema = checks["schema"]
        if not isinstance(response, StructuredResponse):
            raise AssertionError(
                "Response is not a StructuredResponse, cannot check schema"
            )
        if response.parsed is None:
            raise AssertionError("Structured parsing failed: parsed is None")
        if not isinstance(response.parsed, schema):
            raise AssertionError(
                f"Parsed response is {type(response.parsed).__name__}, "
                f"expected {schema.__name__}"
            )


def assert_tool_called(
    agent: Agent,
    tool_name: str,
    *,
    times: int | None = None,
    session_id: str | None = None,
) -> None:
    """Assert that a tool was called during agent execution."""
    state = agent._sessions.get(session_id)
    if state is None:
        raise AssertionError(f"No session found for session_id={session_id!r}")

    count = 0
    for msg in state.messages:
        if msg.role == "assistant" and msg.tool_calls:
            for tc in msg.tool_calls:
                if tc.name == tool_name:
                    count += 1

    if count == 0:
        raise AssertionError(f"Tool '{tool_name}' was never called")

    if times is not None and count != times:
        raise AssertionError(
            f"Tool '{tool_name}' called {count} times, expected {times}"
        )


def assert_tool_not_called(
    agent: Agent,
    tool_name: str,
    *,
    session_id: str | None = None,
) -> None:
    """Assert that a tool was NOT called during agent execution."""
    state = agent._sessions.get(session_id)
    if state is None:
        return  # No session = no calls

    for msg in state.messages:
        if msg.role == "assistant" and msg.tool_calls:
            for tc in msg.tool_calls:
                if tc.name == tool_name:
                    raise AssertionError(f"Tool '{tool_name}' was called but should not have been")


@contextlib.contextmanager
def assert_guardrail_blocked(match: str | None = None):
    """Context manager asserting that a GuardrailError is raised."""
    try:
        yield
    except GuardrailError as exc:
        if match is not None and match not in str(exc):
            raise AssertionError(
                f"GuardrailError raised but message did not match '{match}': {exc}"
            ) from exc
        return

    raise AssertionError("GuardrailError was not raised")
