from __future__ import annotations

from typing import Any

from yuxi.utils.logging_config import logger

_UNSAFE_TRANSPORT_FIELDS = frozenset(
    {
        "agent",
        "cert",
        "cert_file",
        "key",
        "key_file",
        "dispatcher",
        "proxy",
        "session",
    }
)

_AUTH_RESPONSE_MAX_BYTES = 1 * 1024 * 1024


def sanitize_google_auth_init(kwargs: dict[str, Any]) -> dict[str, Any]:
    cleaned: dict[str, Any] = {}
    removed: list[str] = []
    for key, value in kwargs.items():
        if key in _UNSAFE_TRANSPORT_FIELDS:
            removed.append(key)
            continue
        cleaned[key] = value
    if removed:
        logger.warning(f"SSRF guard: removed unsafe transport fields: {removed}")
    return cleaned


def sanitize_request_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
    cleaned: dict[str, Any] = {}
    removed: list[str] = []
    for key, value in kwargs.items():
        if key in _UNSAFE_TRANSPORT_FIELDS:
            removed.append(key)
            continue
        cleaned[key] = value
    if removed:
        logger.warning(f"SSRF guard: removed unsafe request fields: {removed}")
    return cleaned


def apply_ssrf_guard() -> dict[str, Any]:
    return {
        "agent": None,
        "cert": None,
        "cert_file": None,
        "key": None,
        "key_file": None,
    }


def check_auth_response_size(data: bytes) -> bytes:
    if len(data) > _AUTH_RESPONSE_MAX_BYTES:
        logger.warning(
            f"Google Auth response exceeds size limit: {len(data)} bytes > {_AUTH_RESPONSE_MAX_BYTES} bytes, truncating"
        )
        return data[:_AUTH_RESPONSE_MAX_BYTES]
    return data


def check_auth_response_text(text: str) -> str:
    data = text.encode("utf-8", errors="replace")
    truncated = check_auth_response_size(data)
    return truncated.decode("utf-8", errors="replace")