"""数据源代理：服务端替浏览器请求第三方接口

两个入口共用 fetch_json：编辑态由登录用户传 URL；线上态只传数据源 ID，配置取自已发布快照
防护：只允许 http/https；DNS 解析后每个地址都必须是公网地址，并把连接钉在已检查过的地址上
（不再让 httpx 二次解析，防 DNS 重绑定）；不读环境代理；重定向逐跳重新检查，跨域时丢掉调用方的密钥头；
总超时 5 秒，响应体上限 2 MB，按解压后的字节计，解压时带输出上限（防压缩炸弹）
"""

import asyncio
import ipaddress
import json
import re
import socket
import time
import zlib
from collections.abc import Awaitable, Callable
from typing import Any, Literal
from urllib.parse import urljoin

import httpx

from app.core.exceptions import BizException

TIMEOUT = 5.0
MAX_BYTES = 2 * 1024 * 1024
MAX_REDIRECTS = 3
REDIRECT_STATUS = {301, 302, 303, 307, 308}
# 这些头由 httpx 或连接本身决定，不让调用方覆盖
FORBIDDEN_HEADERS = {
    "host",
    "content-length",
    "connection",
    "transfer-encoding",
    "upgrade",
    "te",
    "keep-alive",
    "proxy-authorization",
    "proxy-connection",
    "accept-encoding",
}
# 只接受自己能带上限解压的编码；httpx 默认还会协商 br / zstd，那些解码器没有输出上限
ACCEPT_ENCODING = "gzip, deflate"
# is_global 会放过的几段：NAT64（可以嵌入任意 IPv4）、IPv4 兼容地址、已废弃的站点本地地址
EXTRA_BLOCKED = [
    ipaddress.ip_network(net)
    for net in ("64:ff9b::/96", "64:ff9b:1::/48", "::/96", "fec0::/10")
]
# 十进制、十六进制、省略写法的 IPv4（2130706433、0x7f.1、127.1）：不是标准写法，
# 各平台的解析器解释不一，一律当作内网地址拒绝，不交给 DNS
NUMERIC_LABEL = re.compile(r"0x[0-9a-f]+|[0-9]+", re.IGNORECASE)

Method = Literal["GET", "POST"]


def _blocked() -> BizException:
    return BizException(422, "不允许访问内网或本机地址")


def _is_public(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
    # IPv4 映射的 IPv6 地址按内嵌的 IPv4 判断
    if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
        address = address.ipv4_mapped
    return (
        address.is_global
        and not address.is_multicast
        and not any(address in net for net in EXTRA_BLOCKED)
    )


def _checked(
    address: ipaddress.IPv4Address | ipaddress.IPv6Address,
) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
    if not _is_public(address):
        raise _blocked()
    return address


async def _resolve(host: str, port: int) -> str:
    """解析主机名，所有地址都是公网才放行，返回第一个地址"""
    host = host.rstrip(".")
    try:
        literal = ipaddress.ip_address(host)
    except ValueError:
        pass
    else:
        return str(_checked(literal))
    if host.lower() == "localhost" or host.lower().endswith(".localhost"):
        raise _blocked()
    if all(NUMERIC_LABEL.fullmatch(label) for label in host.split(".")):
        raise _blocked()
    try:
        infos = await asyncio.get_running_loop().getaddrinfo(
            host, port, type=socket.SOCK_STREAM
        )
    except socket.gaierror:
        raise BizException(502, "数据源域名无法解析") from None
    addresses = [ipaddress.ip_address(info[4][0].split("%")[0]) for info in infos]
    if not addresses:
        raise BizException(502, "数据源域名无法解析")
    for address in addresses:
        _checked(address)
    return str(addresses[0])


def _clean_headers(headers: dict[str, Any] | None) -> dict[str, str]:
    cleaned = {"Accept": "application/json", "Accept-Encoding": ACCEPT_ENCODING}
    for name, value in (headers or {}).items():
        if name.lower() in FORBIDDEN_HEADERS or not isinstance(
            value, (str, int, float)
        ):
            continue
        cleaned[name] = str(value)
    return cleaned


def _query_value(value: Any) -> Any:
    return value if isinstance(value, (str, int, float, bool)) else json.dumps(value)


def _parse_url(raw: str) -> httpx.URL:
    try:
        url = httpx.URL(raw)
    except httpx.InvalidURL:
        raise BizException(422, "数据源地址不合法") from None
    if url.scheme not in ("http", "https") or not url.host or url.userinfo:
        raise BizException(422, "数据源地址只支持 http / https，且不能带账号密码")
    return url


def _too_large() -> BizException:
    return BizException(502, "数据源响应超过 2 MB")


async def _read_body(response: httpx.Response) -> bytes:
    """读响应体：原始字节和解压后的字节都不超过 MAX_BYTES

    不用 aiter_bytes：它对每个 64 KB 的块整块解压，没有输出上限，一个块就能膨胀到上 GB
    """
    encoding = response.headers.get("content-encoding", "identity").strip().lower()
    if encoding in ("", "identity"):
        decoder = None
    elif encoding in ("gzip", "x-gzip"):
        decoder = zlib.decompressobj(16 + zlib.MAX_WBITS)
    elif encoding == "deflate":
        decoder = zlib.decompressobj()
    else:
        raise BizException(502, "数据源响应的压缩格式不支持")
    declared = response.headers.get("content-length", "")
    if declared.isdigit() and int(declared) > MAX_BYTES:
        raise _too_large()
    chunks: list[bytes] = []
    raw_size = size = 0
    try:
        async for chunk in response.aiter_raw():
            raw_size += len(chunk)
            if raw_size > MAX_BYTES:
                raise _too_large()
            if decoder is None:
                chunks.append(chunk)
                size += len(chunk)
                continue
            pending = chunk
            while pending:
                # max_length 为 0 表示不限，这里恒大于 0
                part = decoder.decompress(pending, MAX_BYTES - size + 1)
                size += len(part)
                if size > MAX_BYTES:
                    raise _too_large()
                chunks.append(part)
                pending = decoder.unconsumed_tail
                if not part:
                    break
        # 压缩流没有正常结束（被截断）也算失败，别把半截数据当成功
        if decoder is not None and not decoder.eof:
            raise BizException(502, "数据源响应解压失败")
    except zlib.error:
        raise BizException(502, "数据源响应解压失败") from None
    return b"".join(chunks)


async def _request(
    client: httpx.AsyncClient,
    method: str,
    url: httpx.URL,
    headers: dict[str, str],
    body: Any,
) -> tuple[int, httpx.Headers, bytes]:
    port = url.port or (443 if url.scheme == "https" else 80)
    pinned = url.copy_with(host=await _resolve(url.host, port))
    request = client.build_request(
        method,
        pinned,
        headers={**headers, "Host": url.netloc.decode()},
        json=body if method == "POST" else None,
        # 连接用 IP，TLS 的 SNI 和证书校验仍然用原主机名
        extensions={"sni_hostname": url.host},
    )
    response = await client.send(request, stream=True)
    try:
        return response.status_code, response.headers, await _read_body(response)
    finally:
        await response.aclose()


def _origin(url: httpx.URL) -> tuple[str, str, int | None]:
    return url.scheme, url.host, url.port or (443 if url.scheme == "https" else 80)


async def fetch_json(
    url: str,
    method: Method = "GET",
    params: dict[str, Any] | None = None,
    headers: dict[str, Any] | None = None,
) -> dict[str, Any] | list[Any]:
    """GET 的 params 作为查询参数，POST 的 params 作为 JSON 请求体，和前端 fetchData 一致"""
    target = _parse_url(url)
    request_headers = _clean_headers(headers)
    body = None
    if method == "GET":
        if params:
            target = target.copy_merge_params(
                {key: _query_value(value) for key, value in params.items()}
            )
    else:
        body = params or {}
    try:
        async with asyncio.timeout(TIMEOUT):
            # trust_env=False：环境里的代理会绕过上面的地址检查
            # max_keepalive_connections=0：连接按 IP 复用，不关的话同一 IP 上的下一跳会沿用上一跳的 TLS 校验结果
            async with httpx.AsyncClient(
                trust_env=False,
                timeout=TIMEOUT,
                limits=httpx.Limits(max_keepalive_connections=0),
            ) as client:
                for _ in range(MAX_REDIRECTS + 1):
                    status, response_headers, content = await _request(
                        client, method, target, request_headers, body
                    )
                    location = response_headers.get("location")
                    if status not in REDIRECT_STATUS or not location:
                        break
                    previous = target
                    target = _parse_url(urljoin(str(target), location))
                    if _origin(target) != _origin(previous):
                        # 密钥头只发给调用方指定的那个源
                        request_headers = _clean_headers(None)
                    if status not in (307, 308):
                        method, body = "GET", None
                else:
                    raise BizException(502, "数据源重定向次数过多")
    except (TimeoutError, httpx.TimeoutException):
        raise BizException(504, "数据源请求超时") from None
    except httpx.HTTPError:
        raise BizException(502, "数据源请求失败") from None
    if not 200 <= status < 300:
        raise BizException(502, f"数据源返回 HTTP {status}")
    try:
        data = json.loads(content)
    except (ValueError, RecursionError):
        # 2 MB 以内的深层嵌套会打满递归，同样当作非法 JSON
        raise BizException(502, "数据源返回的不是 JSON") from None
    if not isinstance(data, (dict, list)):
        raise BizException(502, "数据源返回的不是 JSON 对象或数组")
    return data


class FetchCache:
    """线上大屏多人同时开、按 interval 轮询：同一个 key 几秒内只真正请求一次，并发的请求共用同一次

    失败不缓存；单实例部署，缓存放进程内存。单条最大 2 MB，条数上限乘出来就是内存上限，所以条数要小
    """

    def __init__(self, ttl: float = 3.0, max_entries: int = 64):
        self._ttl = ttl
        self._max_entries = max_entries
        self._entries: dict[Any, tuple[float, Any]] = {}
        self._inflight: dict[Any, asyncio.Future] = {}

    async def get(self, key: Any, fetch: Callable[[], Awaitable[Any]]) -> Any:
        now = time.monotonic()
        hit = self._entries.get(key)
        if hit is not None and hit[0] > now:
            return hit[1]
        future = self._inflight.get(key)
        if future is None:
            future = asyncio.ensure_future(fetch())
            self._inflight[key] = future
            future.add_done_callback(lambda done: self._finish(key, done))
        # shield：某个请求者断开不能取消共用的那次请求
        return await asyncio.shield(future)

    def _finish(self, key: Any, future: asyncio.Future) -> None:
        self._inflight.pop(key, None)
        if future.cancelled() or future.exception() is not None:
            return
        now = time.monotonic()
        # 每次写入先丢掉过期项；清完还满就不缓存这一条
        self._entries = {k: v for k, v in self._entries.items() if v[0] > now}
        if len(self._entries) < self._max_entries:
            self._entries[key] = (now + self._ttl, future.result())