"""
Test assertions for automation testing.

This module provides assertion helpers that make it easy to write
tests for automations.
"""

import re
from typing import Any

from torivers_sdk.context.progress import LogLevel
from torivers_sdk.testing.sandbox import ExecutionResult


def _resolve_output_key(output_data: dict[str, Any], key: str) -> tuple[bool, Any]:
    """
    Resolve a key from json_data block payloads.
    """
    blocks = output_data.get("_blocks")
    if isinstance(blocks, list):
        for block in blocks:
            if not isinstance(block, dict):
                continue
            if block.get("type") != "json_data":
                continue
            data = block.get("data")
            if isinstance(data, dict) and key in data:
                return True, data[key]

    return False, None


def assert_output_contains(
    result: ExecutionResult,
    key: str,
    message: str | None = None,
) -> None:
    """
    Assert that the output data contains a specific key.

    Args:
        result: Execution result from sandbox
        key: Key that should exist in json_data block payload data
        message: Optional custom error message

    Raises:
        AssertionError: If key not in output_data
    """
    found, _ = _resolve_output_key(result.output_data, key)
    if not found:
        error_msg = message or f"Output data missing key: {key}"
        raise AssertionError(error_msg)


def assert_output_matches(
    result: ExecutionResult,
    key: str,
    expected: Any,
    message: str | None = None,
) -> None:
    """
    Assert that an output value matches expected value.

    Args:
        result: Execution result from sandbox
        key: Key in json_data block payload data
        expected: Expected value
        message: Optional custom error message

    Raises:
        AssertionError: If value doesn't match
    """
    found, actual = _resolve_output_key(result.output_data, key)
    if not found:
        error_msg = message or f"Output data missing key: {key}"
        raise AssertionError(error_msg)
    if actual != expected:
        error_msg = message or f"Output[{key}] = {actual!r}, expected {expected!r}"
        raise AssertionError(error_msg)


def assert_output_type(
    result: ExecutionResult,
    key: str,
    expected_type: type,
    message: str | None = None,
) -> None:
    """
    Assert that an output value is of expected type.

    Args:
        result: Execution result from sandbox
        key: Key in json_data block payload data
        expected_type: Expected type
        message: Optional custom error message

    Raises:
        AssertionError: If value is wrong type
    """
    found, actual = _resolve_output_key(result.output_data, key)
    if not found:
        error_msg = message or f"Output data missing key: {key}"
        raise AssertionError(error_msg)
    if not isinstance(actual, expected_type):
        error_msg = message or (
            f"Output[{key}] is {type(actual).__name__}, "
            f"expected {expected_type.__name__}"
        )
        raise AssertionError(error_msg)


def assert_progress_logged(
    result: ExecutionResult,
    title_pattern: str,
    level: LogLevel | None = None,
    message: str | None = None,
) -> None:
    """
    Assert that a progress entry was logged.

    Args:
        result: Execution result from sandbox
        title_pattern: Regex pattern to match against entry titles
        level: Optional log level to filter by
        message: Optional custom error message

    Raises:
        AssertionError: If no matching entry found
    """
    pattern = re.compile(title_pattern)

    for entry in result.progress_entries:
        if level is not None and entry.level != level:
            continue
        if pattern.search(entry.title):
            return

    error_msg = message or f"No progress entry matching '{title_pattern}'"
    if level:
        error_msg += f" with level {level.value}"
    raise AssertionError(error_msg)


def assert_no_errors(
    result: ExecutionResult,
    message: str | None = None,
) -> None:
    """
    Assert that execution completed without errors.

    Args:
        result: Execution result from sandbox
        message: Optional custom error message

    Raises:
        AssertionError: If errors occurred
    """
    if result.errors:
        error_msg = message or f"Execution had errors: {result.errors}"
        raise AssertionError(error_msg)


def assert_success(
    result: ExecutionResult,
    message: str | None = None,
) -> None:
    """
    Assert that execution was successful.

    Args:
        result: Execution result from sandbox
        message: Optional custom error message

    Raises:
        AssertionError: If execution was not successful
    """
    if not result.success:
        error_msg = message or f"Execution failed: {result.errors}"
        raise AssertionError(error_msg)


def assert_execution_time(
    result: ExecutionResult,
    max_seconds: float,
    message: str | None = None,
) -> None:
    """
    Assert that execution completed within time limit.

    Args:
        result: Execution result from sandbox
        max_seconds: Maximum allowed execution time
        message: Optional custom error message

    Raises:
        AssertionError: If execution took too long
    """
    if result.execution_time_seconds > max_seconds:
        error_msg = message or (
            f"Execution took {result.execution_time_seconds:.2f}s, "
            f"expected max {max_seconds:.2f}s"
        )
        raise AssertionError(error_msg)


def assert_within_limits(
    result: ExecutionResult,
    max_duration: float | None = None,
    max_memory: float | None = None,
    max_tokens: int | None = None,
    message: str | None = None,
) -> None:
    """
    Assert that execution was within resource limits.

    This is a comprehensive check that can verify multiple resource
    constraints in a single assertion.

    Args:
        result: Execution result from sandbox
        max_duration: Maximum allowed execution time in seconds
        max_memory: Maximum allowed peak memory in megabytes
        max_tokens: Maximum allowed tokens used
        message: Optional custom error message prefix

    Raises:
        AssertionError: If any resource limit was exceeded

    Example:
        # Check execution stayed within production-like limits
        assert_within_limits(
            result,
            max_duration=60.0,  # 60 seconds
            max_memory=512.0,   # 512 MB
            max_tokens=10000,   # 10k tokens
        )
    """
    violations = []

    if max_duration is not None:
        if result.execution_time_seconds > max_duration:
            violations.append(
                f"duration {result.execution_time_seconds:.2f}s > {max_duration:.2f}s"
            )

    if max_memory is not None:
        # Check if result has memory tracking
        memory_mb = getattr(result, "peak_memory_mb", 0.0)
        if memory_mb > max_memory:
            violations.append(f"memory {memory_mb:.2f}MB > {max_memory:.2f}MB")

    if max_tokens is not None:
        # Check if result has token tracking
        tokens = getattr(result, "total_tokens_used", 0)
        if tokens > max_tokens:
            violations.append(f"tokens {tokens} > {max_tokens}")

    if violations:
        prefix = f"{message}: " if message else "Resource limits exceeded: "
        error_msg = prefix + ", ".join(violations)
        raise AssertionError(error_msg)


def assert_tokens_within_limit(
    result: ExecutionResult,
    max_tokens: int,
    message: str | None = None,
) -> None:
    """
    Assert that token usage is within limit.

    Args:
        result: Execution result from sandbox
        max_tokens: Maximum allowed tokens
        message: Optional custom error message

    Raises:
        AssertionError: If token usage exceeds limit
    """
    tokens = getattr(result, "total_tokens_used", 0)
    if tokens > max_tokens:
        error_msg = message or f"Token usage {tokens} exceeded limit of {max_tokens}"
        raise AssertionError(error_msg)


def assert_memory_within_limit(
    result: ExecutionResult,
    max_memory_mb: float,
    message: str | None = None,
) -> None:
    """
    Assert that peak memory usage is within limit.

    Args:
        result: Execution result from sandbox
        max_memory_mb: Maximum allowed memory in megabytes
        message: Optional custom error message

    Raises:
        AssertionError: If memory usage exceeds limit
    """
    memory_mb = getattr(result, "peak_memory_mb", 0.0)
    if memory_mb > max_memory_mb:
        error_msg = message or (
            f"Peak memory {memory_mb:.2f}MB exceeded limit of {max_memory_mb:.2f}MB"
        )
        raise AssertionError(error_msg)
