import asyncio
from typing import Any

import aiohttp

from astraflow.raas.utils import logging

DEFAULT_RETRIES = 1
DEFAULT_REQUEST_TIMEOUT = 3600


logger = logging.getLogger(__file__)


def get_default_connector():
    return aiohttp.TCPConnector(limit=0, use_dns_cache=False, force_close=True)


async def arequest_with_retry(
    addr: str,
    endpoint: str,
    payload: dict[str, Any] | None = None,
    session: aiohttp.ClientSession | None = None,
    method: str = "POST",
    max_retries: int | None = None,
    timeout: float | None = None,
    retry_delay: float = 1.0,
    verbose=False,
) -> dict:
    if timeout is None:
        timeout = DEFAULT_REQUEST_TIMEOUT
    last_exception = None
    max_retries = max_retries or DEFAULT_RETRIES
    base_url = f"http://{addr}"
    url = f"{base_url}{endpoint}"

    timeo = aiohttp.ClientTimeout(
        total=timeout,
        sock_connect=timeout,
        connect=timeout,
    )
    if session is None:
        _session = aiohttp.ClientSession(
            timeout=timeo,
            read_bufsize=1024 * 1024 * 10,
            connector=get_default_connector(),
        )
    else:
        _session = session

    for attempt in range(max_retries):
        try:
            if verbose:
                logger.info("enter client session, start sending requests")
            if method.upper() == "GET":
                ctx = _session.get(url, timeout=timeo)
            elif method.upper() == "POST":
                ctx = _session.post(url, json=payload, timeout=timeo)
            elif method.upper() == "PUT":
                ctx = _session.put(url, json=payload, timeout=timeo)
            elif method.upper() == "DELETE":
                ctx = _session.delete(url, timeout=timeo)
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")
            async with ctx as response:
                if verbose:
                    logger.info("http requests return")
                if response.status >= 400:
                    body = await response.text()
                    logger.warning(
                        "HTTP %d from %s%s: %s",
                        response.status,
                        addr,
                        endpoint,
                        body[:500],
                    )
                response.raise_for_status()
                try:
                    res = await response.json()
                except aiohttp.ContentTypeError:
                    # Some SGLang endpoints (e.g. /flush_cache) return
                    # plain text instead of JSON.
                    text = await response.text()
                    res = {"text": text}
                if verbose:
                    logger.info("get http result")
                if session is None:
                    await _session.close()
                return res
        except (
            aiohttp.ClientError,
            aiohttp.ClientResponseError,
            asyncio.TimeoutError,
        ) as e:
            if isinstance(e, asyncio.TimeoutError):
                logger.warning(
                    "HTTP request to %s%s timed out after %.2fs (attempt %d/%d)",
                    addr,
                    endpoint,
                    timeout,
                    attempt + 1,
                    max_retries,
                )
            else:
                logger.warning(
                    "HTTP request to %s%s failed with %s: %s (attempt %d/%d)",
                    addr,
                    endpoint,
                    e.__class__.__name__,
                    str(e),
                    attempt + 1,
                    max_retries,
                )
            last_exception = e
            if attempt < max_retries - 1:
                await asyncio.sleep(retry_delay)
            continue
    if session is None:
        await _session.close()
    # Summarise the payload so that huge fields (e.g. input_ids) don't
    # blow up log lines.
    def _summarise(obj, max_len=200):
        s = str(obj)
        if len(s) <= max_len:
            return s
        return s[:max_len] + f"... (truncated, total {len(s)} chars)"

    raise RuntimeError(
        f"Failed after {max_retries} retries each. "
        f"Payload: {_summarise(payload)}. Addr: {addr}. Endpoint: {endpoint}. "
        f"Last error: {repr(last_exception)}"
    )
