from urllib.parse import urljoin, urlparse

import httpx

from yuxi.utils import logger
from yuxi.utils.net_security import SSRFValidator, SSRFPolicy, log_ssrf_block
from yuxi.knowledge.utils.url_validator import is_url_parsing_enabled

# 最大允许下载大小 (例如 10MB)
MAX_DOWNLOAD_SIZE = 10 * 1024 * 1024
# 允许的 Content-Type
ALLOWED_CONTENT_TYPES = ["text/html", "application/xhtml+xml"]

_validator: SSRFValidator | None = None


def _get_validator() -> SSRFValidator:
    """获取 SSRF 验证器单例（从环境变量构建策略）。"""
    global _validator
    if _validator is None:
        import os

        whitelist = [d.strip() for d in os.environ.get("YUXI_URL_WHITELIST", "").split(",") if d.strip()]
        policy = SSRFPolicy(allowed_domains=whitelist)
        _validator = SSRFValidator(policy)
    return _validator


async def fetch_url_content(url: str, max_size: int = MAX_DOWNLOAD_SIZE) -> tuple[bytes, str]:
    """
    Fetch URL content with security checks (SSRF protection, size limit, content type).

    Args:
        url: The URL to fetch.
        max_size: Maximum allowed size in bytes.

    Returns:
        tuple: (content_bytes, final_url)

    Raises:
        ValueError: If validation fails or download error occurs.
    """
    if not is_url_parsing_enabled():
        raise ValueError("URL parsing feature is disabled")

    validator = _get_validator()

    # 初始 SSRF 校验
    allowed, reason = await validator.validate_url(url)
    if not allowed:
        parsed = urlparse(url)
        log_ssrf_block(parsed.hostname or url, reason, context="url-fetch")
        raise ValueError(f"请求被安全策略阻止: {reason}")

    current_url = url
    redirect_count = 0
    max_redirects = 5

    # 手动处理重定向，每次重定向重新校验
    try:
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client:
            while True:
                logger.info(f"Fetching URL: {current_url}")

                headers = {
                    "User-Agent": (
                        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                        "AppleWebKit/537.36 (KHTML, like Gecko) "
                        "Chrome/91.0.4472.124 "
                        "Safari/537.36"
                    )
                }

                async with client.stream("GET", current_url, headers=headers) as response:
                    # Handle Redirects
                    if response.status_code in (301, 302, 303, 307, 308):
                        if redirect_count >= max_redirects:
                            raise ValueError("Too many redirects")

                        redirect_count += 1
                        location = response.headers.get("Location")
                        if not location:
                            raise ValueError("Redirect response missing Location header")

                        current_url = urljoin(current_url, location)

                        # 重定向 SSRF 校验
                        allowed, reason = await validator.validate_url(current_url)
                        if not allowed:
                            parsed = urlparse(current_url)
                            log_ssrf_block(parsed.hostname or current_url, reason, context="url-fetch-redirect")
                            raise ValueError(f"重定向被安全策略阻止: {reason}")

                        continue  # Start new request

                    response.raise_for_status()

                    # Check Content-Type
                    content_type = response.headers.get("Content-Type", "").lower()
                    if not any(allowed in content_type for allowed in ALLOWED_CONTENT_TYPES):
                        raise ValueError(f"Unsupported Content-Type: {content_type}. Only HTML is supported.")

                    # Download content with size limit
                    content = bytearray()
                    async for chunk in response.aiter_bytes():
                        content.extend(chunk)
                        if len(content) > max_size:
                            raise ValueError(f"Content size exceeds limit of {max_size} bytes")

                    return bytes(content), current_url

    except httpx.HTTPError as e:
        logger.error(f"HTTP error fetching {url}: {e}")
        raise ValueError(f"Failed to fetch URL: {e}")
    except Exception as e:
        logger.error(f"Error fetching {url}: {e}")
        raise ValueError(f"Error fetching URL: {str(e)}")