# chuk_mcp/protocol/messages/send_message.py
import logging
import uuid
from typing import Any, Dict, Optional, Union, Callable, Awaitable, List

import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream

# chuk_mcp imports
from chuk_mcp.protocol.messages.json_rpc_message import JSONRPCMessage
from chuk_mcp.protocol.types.errors import get_error_message, is_retryable_error
from chuk_mcp.protocol.types.errors import RetryableError, NonRetryableError

###############################################################################
# Cancellation Support
###############################################################################


class CancellationToken:
    """Token to enable request cancellation."""

    def __init__(self) -> None:
        self._cancelled = False
        self._callbacks: List[Callable] = []

    def cancel(self):
        """Mark this token as cancelled."""
        self._cancelled = True
        for callback in self._callbacks:
            try:
                callback()
            except Exception as e:
                logging.error(f"Error in cancellation callback: {e}")

    @property
    def is_cancelled(self) -> bool:
        """Check if cancellation was requested."""
        return self._cancelled

    def add_callback(self, callback: Callable) -> None:
        """Add a callback to be called on cancellation."""
        self._callbacks.append(callback)
        if self._cancelled:
            callback()


class CancelledError(Exception):
    """Raised when a request is cancelled via CancellationToken."""

    pass


###############################################################################
# Public API - Enhanced JSON-RPC send/receive
###############################################################################
async def send_message(
    read_stream: MemoryObjectReceiveStream,
    write_stream: MemoryObjectSendStream,
    method: str,
    params: Optional[Dict[str, Any]] = None,
    *,
    timeout: float = 5.0,
    message_id: Optional[str] = None,
    cancellation_token: Optional[CancellationToken] = None,
    progress_callback: Optional[
        Callable[[float, Optional[float], Optional[str]], Awaitable[None]]
    ] = None,
) -> Union[Dict[str, Any], Any]:
    """Send a JSON-RPC 2.0 request message and await the matching response.

    This is a clean protocol layer - errors are raised immediately without retries.
    Higher layers should implement retry logic if needed.

    Args:
        read_stream: Stream to receive JSONRPCMessage objects.
        write_stream: Stream to send JSONRPCMessage objects.
        method: The JSON-RPC method name.
        params: Parameters dict for the call.
        timeout: Timeout in seconds for the response.
        message_id: Custom request ID; autogenerated if None.
        cancellation_token: Token to enable request cancellation
        progress_callback: Callback for progress updates (progress, total, message)

    Returns:
        The 'result' field of the JSON-RPC response, or full response dict.

    Raises:
        TimeoutError: If no matching response arrives in time.
        RetryableError: For retryable JSON-RPC errors.
        NonRetryableError: For terminal JSON-RPC errors.
        CancelledError: If the request was cancelled via the token.
    """
    # Generate progress token if callback provided
    progress_token = None
    if progress_callback:
        progress_token = str(uuid.uuid4())
        if params is None:
            params = {}
        # Add progress token to _meta
        if "_meta" not in params:
            params["_meta"] = {}
        params["_meta"]["progressToken"] = progress_token

    # Prepare request - import the actual function
    from chuk_mcp.protocol.messages.json_rpc_message import create_request

    req_id = message_id or str(uuid.uuid4())
    message = create_request(method=method, params=params, id=req_id)

    # Track if we've sent a cancellation
    cancellation_sent = False

    async def check_and_send_cancellation():
        nonlocal cancellation_sent
        if (
            cancellation_token
            and cancellation_token.is_cancelled
            and not cancellation_sent
        ):
            cancellation_sent = True
            try:
                # Import here to avoid circular dependency
                from chuk_mcp.protocol.messages.notifications import (
                    send_cancelled_notification,
                )

                await send_cancelled_notification(
                    write_stream, req_id, "Cancelled by client"
                )
            except Exception as e:
                logging.error(f"Failed to send cancellation notification: {e}")
            raise CancelledError(f"Request {req_id} was cancelled")

    # Register cancellation callback if token provided
    if cancellation_token:
        cancellation_token.add_callback(
            lambda: None
        )  # Just to trigger check on next iteration

    # Check for cancellation before sending
    if cancellation_token:
        await check_and_send_cancellation()

    logging.debug("[send_message] sending %s", method)
    await write_stream.send(message)

    with anyio.fail_after(timeout):
        return await _await_response(
            read_stream,
            req_id,
            cancellation_check=check_and_send_cancellation
            if cancellation_token
            else None,
            progress_token=progress_token,
            progress_callback=progress_callback,
        )


###############################################################################
# Helper - iterative receive with sub-timeout
###############################################################################
async def _await_response(
    read_stream: MemoryObjectReceiveStream,
    req_id: str,
    sub_timeout: float = 0.5,
    cancellation_check: Optional[Callable[[], Awaitable[None]]] = None,
    progress_token: Optional[str] = None,
    progress_callback: Optional[
        Callable[[float, Optional[float], Optional[str]], Awaitable[None]]
    ] = None,
) -> Union[Dict[str, Any], Any]:
    """Loop until a response matching req_id arrives, handling progress notifications."""
    while True:
        # Check for cancellation
        if cancellation_check:
            await cancellation_check()

        try:
            with anyio.fail_after(sub_timeout):
                msg = await read_stream.receive()
        except TimeoutError:
            continue  # let outer timer count down

        # Handle progress notifications if we're tracking progress
        # Type narrowing - only requests/notifications have method
        msg_method = getattr(msg, "method", None)
        if (
            progress_token
            and progress_callback
            and msg_method == "notifications/progress"
        ):
            params = getattr(msg, "params", None) or {}
            if params.get("progressToken") == progress_token:
                try:
                    await progress_callback(
                        params.get("progress", 0),
                        params.get("total"),
                        params.get("message"),
                    )
                except Exception as e:
                    logging.error(f"Error in progress callback: {e}")
                continue  # Keep waiting for the actual response

        # Filter by matching ID - type narrowing
        msg_id = getattr(msg, "id", None)
        if msg_id != req_id:
            logging.debug("[send_message] skip unmatched id=%s", msg_id)
            continue

        # For mypy - msg must be a single message, not a list
        if isinstance(msg, list):
            continue

        logging.debug("[send_message] matched response: %s", msg.model_dump())
        return _process_response(msg)


###############################################################################
# Common - response error handling
###############################################################################


def _process_response(resp: JSONRPCMessage) -> Union[Dict[str, Any], Any]:
    # Type narrowing for union - only JSONRPCError and JSONRPCResponse have these attributes
    error = getattr(resp, "error", None)
    if error is not None:
        code = error.get("code", -32603)
        msg = (
            f"JSON-RPC Error: {error.get('message', get_error_message(code))}"
            f" (code: {code})"
        )
        if is_retryable_error(code):
            raise RetryableError(msg, code)
        raise NonRetryableError(msg, code)
    # Successful response
    result = getattr(resp, "result", None)
    if result is not None:
        return result
    # For lists and other types, return model_dump if available
    if hasattr(resp, "model_dump"):
        return resp.model_dump()  # type: ignore[union-attr]
    return resp  # type: ignore[return-value]
