TelegramBotAdmin/wbb/services/directory_geocoding.py

from __future__ import annotations

import asyncio
from contextlib import suppress
from typing import Any

from aiohttp import ClientError, ClientSession, ClientTimeout

from wbb import log
from wbb.utils.dbdirectory import (
    claim_geocode_job,
    complete_geocode_job,
    fail_geocode_job,
    get_directory_settings,
)

NOMINATIM_ATTRIBUTION = "© OpenStreetMap contributors"
NOMINATIM_ATTRIBUTION_URL = "https://www.openstreetmap.org/copyright"

_worker_task: asyncio.Task[None] | None = None
_stop_event: asyncio.Event | None = None


def parse_nominatim_region(payload: dict[str, Any]) -> dict[str, str]:
    address = payload.get("address") if isinstance(payload.get("address"), dict) else {}
    province = str(address.get("state") or address.get("province") or "").strip()
    city = str(
        address.get("city")
        or address.get("municipality")
        or address.get("town")
        or address.get("county")
        or ""
    ).strip()
    district = str(
        address.get("city_district")
        or address.get("district")
        or address.get("borough")
        or address.get("suburb")
        or address.get("county")
        or ""
    ).strip()
    values = list(dict.fromkeys(value for value in (province, city, district) if value))
    return {
        "province": province,
        "city": city,
        "district": district,
        "label": " ".join(values) if values else "区域暂不可用",
        "provider": "nominatim",
    }


async def reverse_geocode(
    *,
    endpoint: str,
    contact: str,
    user_agent: str,
    longitude: float,
    latitude: float,
) -> dict[str, str]:
    headers = {
        "User-Agent": f"{user_agent} ({contact})",
        "Accept": "application/json",
        "Accept-Language": "zh-CN,zh;q=0.9",
    }
    params = {
        "format": "jsonv2",
        "lat": repr(float(latitude)),
        "lon": repr(float(longitude)),
        "zoom": "12",
        "addressdetails": "1",
        "accept-language": "zh-CN",
    }
    timeout = ClientTimeout(total=12)
    async with ClientSession(timeout=timeout, headers=headers) as session:
        async with session.get(endpoint, params=params) as response:
            if response.status == 429:
                retry_after = int(response.headers.get("Retry-After") or 60)
                raise NominatimRateLimited(retry_after)
            if response.status >= 400:
                text = await response.text()
                raise RuntimeError(
                    f"Nominatim 返回 HTTP {response.status}: {text[:200]}"
                )
            payload = await response.json()
    if not isinstance(payload, dict):
        raise RuntimeError("Nominatim 返回了无效数据。")
    return parse_nominatim_region(payload)


class NominatimRateLimited(RuntimeError):
    def __init__(self, retry_after_seconds: int):
        super().__init__("Nominatim 请求已被限流。")
        self.retry_after_seconds = max(1, int(retry_after_seconds))


async def _geocode_worker(stop_event: asyncio.Event) -> None:
    last_request_at = 0.0
    while not stop_event.is_set():
        settings = await get_directory_settings()
        enabled = bool(settings.get("nominatim_enabled"))
        contact = str(settings.get("nominatim_contact") or "").strip()
        if not enabled or not contact:
            with suppress(TimeoutError):
                await asyncio.wait_for(stop_event.wait(), timeout=2)
            continue
        job = await claim_geocode_job()
        if not job:
            with suppress(TimeoutError):
                await asyncio.wait_for(stop_event.wait(), timeout=1)
            continue
        loop = asyncio.get_running_loop()
        wait_seconds = 1.1 - (loop.time() - last_request_at)
        if wait_seconds > 0:
            with suppress(TimeoutError):
                await asyncio.wait_for(stop_event.wait(), timeout=wait_seconds)
            if stop_event.is_set():
                return
        try:
            last_request_at = loop.time()
            region = await reverse_geocode(
                endpoint=str(settings["nominatim_endpoint"]),
                contact=contact,
                user_agent=str(settings["nominatim_user_agent"]),
                longitude=float(job["longitude"]),
                latitude=float(job["latitude"]),
            )
            await complete_geocode_job(job=job, region=region)
        except NominatimRateLimited as exc:
            await fail_geocode_job(
                job=job,
                error=str(exc),
                retry_after_seconds=exc.retry_after_seconds,
            )
        except (ClientError, TimeoutError, ValueError, RuntimeError) as exc:
            log.error(f"Nominatim 区域解析失败：{exc}")
            await fail_geocode_job(job=job, error=str(exc), retry_after_seconds=60)
        except Exception as exc:
            log.error(f"Nominatim 区域解析出现未预期错误：{exc}")
            await fail_geocode_job(job=job, error="区域解析服务异常。")


async def start_directory_geocoder() -> None:
    global _worker_task, _stop_event
    if _worker_task and not _worker_task.done():
        return
    _stop_event = asyncio.Event()
    _worker_task = asyncio.create_task(
        _geocode_worker(_stop_event),
        name="directory-geocoder",
    )


async def stop_directory_geocoder() -> None:
    global _worker_task, _stop_event
    if _stop_event:
        _stop_event.set()
    if _worker_task:
        with suppress(asyncio.CancelledError):
            await asyncio.wait_for(_worker_task, timeout=3)
    _worker_task = None
    _stop_event = None