"""MCP-specific assertion functions for the test harness.

Each assertion returns a tuple of ``(passed, reason)`` where *passed* is a
bool and *reason* is a human-readable explanation (useful for both pass and
fail reporting).
"""

from __future__ import annotations

import json
from typing import Any

import jsonschema

from hatchdx.harness.simulator import ToolCallResult

AssertionResult = tuple[bool, str]


# ---------------------------------------------------------------------------
# Individual assertions
# ---------------------------------------------------------------------------


def assert_no_error(result: ToolCallResult) -> AssertionResult:
    """Assert the tool call did not return an error."""
    if result.is_error:
        text = _extract_text(result)
        return False, f"Expected success but got error: {text[:200]}"
    return True, "Response is not an error"


def assert_is_error(result: ToolCallResult) -> AssertionResult:
    """Assert the tool call returned an error."""
    if not result.is_error:
        return False, "Expected an error response but got success"
    return True, "Response is an error as expected"


def assert_content_contains(result: ToolCallResult, text: str) -> AssertionResult:
    """Assert the response text content contains *text*."""
    full_text = _extract_text(result)
    if text in full_text:
        return True, f"Content contains '{text}'"
    return False, f"Expected content to contain '{text}', got: {full_text[:200]}"


def assert_content_not_contains(result: ToolCallResult, text: str) -> AssertionResult:
    """Assert the response text content does NOT contain *text*."""
    full_text = _extract_text(result)
    if text not in full_text:
        return True, f"Content does not contain '{text}'"
    return False, f"Expected content to NOT contain '{text}', but it was found"


def assert_schema(result: ToolCallResult, schema: dict[str, Any]) -> AssertionResult:
    """Validate the text content as JSON against a JSON Schema."""
    full_text = _extract_text(result)
    if not full_text:
        return False, "No text content to validate against schema"

    try:
        data = json.loads(full_text)
    except json.JSONDecodeError:
        return False, f"Content is not valid JSON: {full_text[:200]}"

    try:
        jsonschema.validate(instance=data, schema=schema)
    except jsonschema.ValidationError as exc:
        return False, f"Schema validation failed: {exc.message}"

    return True, "Content matches JSON Schema"


def assert_latency_under(latency_ms: float, max_ms: int) -> AssertionResult:
    """Assert the tool call completed within *max_ms* milliseconds."""
    if latency_ms <= max_ms:
        return True, f"Latency {latency_ms:.1f}ms is under {max_ms}ms limit"
    return False, f"Latency {latency_ms:.1f}ms exceeds limit of {max_ms}ms"


def assert_valid_mcp_response(result: ToolCallResult) -> AssertionResult:
    """Assert the response is structurally valid per MCP protocol.

    Checks:
    - ``content`` is a list
    - Each block is a dict with a ``type`` field
    - Text blocks have a ``text`` field
    - Image blocks have ``data`` and ``mimeType`` fields
    """
    if not isinstance(result.content, list):
        return (
            False,
            f"Expected content to be a list, got {type(result.content).__name__}",
        )

    for i, block in enumerate(result.content):
        if not isinstance(block, dict):
            return False, f"Content block {i} is not a mapping"
        if "type" not in block:
            return False, f"Content block {i} is missing 'type' field"

        block_type = block["type"]
        if block_type == "text" and "text" not in block:
            return False, f"Text content block {i} is missing 'text' field"
        if block_type == "image" and ("data" not in block or "mimeType" not in block):
            return (
                False,
                f"Image content block {i} is missing 'data' or 'mimeType' field",
            )

    return True, "Response is a valid MCP tool result"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _extract_text(result: ToolCallResult) -> str:
    """Join all text content blocks from a ToolCallResult."""
    parts: list[str] = []
    for block in result.content:
        if isinstance(block, dict) and block.get("type") == "text":
            parts.append(block.get("text", ""))
    return "\n".join(parts)
