# 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):
        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,
    retries: int = 3,
    retry_delay: float = 2.0,
    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.

    Uses a short per-iteration sub-timeout to keep the outer timeout accurate
    even if the server stalls.

    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: Total timeout for a single attempt.
        message_id: Custom request ID; autogenerated if None.
        retries: Number of attempts before giving up.
        retry_delay: Delay between attempts.
        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
    req_id = message_id or str(uuid.uuid4())
    message = JSONRPCMessage(id=req_id, method=method, params=params)

    # 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

    last_exc: Optional[Exception] = None
    for attempt in range(1, retries + 1):
        try:
            # Check for cancellation before each attempt
            if cancellation_token:
                await check_and_send_cancellation()
            
            logging.debug("[send_message] attempt %s/%s -> %s", attempt, retries, 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
                )

        except CancelledError:
            # Cancellation is terminal, don't retry
            raise
        except NonRetryableError:
            # Permanent error, abort retries
            raise
        except RetryableError as exc:
            last_exc = exc
            logging.error("[send_message] retryable error: %s", exc)
        except TimeoutError as exc:
            last_exc = exc
            logging.error("[send_message] timeout for %s on attempt %s/%s", method, attempt, retries)
        except Exception as exc:
            last_exc = exc
            logging.error("[send_message] unexpected error: %s", exc)

        # Backoff before retrying
        if attempt < retries:
            await anyio.sleep(retry_delay)

    # All retries exhausted
    assert last_exc is not None
    raise last_exc

###############################################################################
# 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
        if progress_token and progress_callback and msg.method == "notifications/progress":
            params = msg.params 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
        if msg.id != req_id:
            logging.debug("[send_message] skip unmatched id=%s", msg.id)
            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]:
    if resp.error is not None:
        code = resp.error.get("code", -32603)
        msg = (
            f"JSON-RPC Error: {resp.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
    return resp.result if resp.result is not None else resp.model_dump()