# chuk_mcp/mcp_client/messages/send_message.py
import logging
import uuid
from typing import Any, Dict, Optional, Union

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

# chuk_mcp imports
from chuk_mcp.mcp_client.messages.json_rpc_message import JSONRPCMessage
from chuk_mcp.mcp_client.messages.error_codes import get_error_message, is_retryable_error
from chuk_mcp.mcp_client.messages.exceptions import RetryableError, NonRetryableError

###############################################################################
# Public API – legacy 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,
) -> 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.

    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.
    """
    # Prepare request
    req_id = message_id or str(uuid.uuid4())
    message = JSONRPCMessage(id=req_id, method=method, params=params)

    last_exc: Optional[Exception] = None
    for attempt in range(1, retries + 1):
        try:
            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)

        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,
) -> Union[Dict[str, Any], Any]:
    """Loop until a response matching req_id arrives, using short time slices."""
    while True:
        try:
            with anyio.fail_after(sub_timeout):
                resp = await read_stream.receive()
        except TimeoutError:
            continue  # let outer timer count down

        # Filter by matching ID
        if resp.id != req_id:
            logging.debug("[send_message] skip unmatched id=%s", resp.id)
            continue

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

###############################################################################
# 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()
