sub-provider/app/services/subscriptions.py

from __future__ import annotations

import re
from typing import Any, Iterable

import httpx
import yaml

from app.config import get_settings
from app.models import FetchResult, ProviderDocument, SourceConfig, SourceSnapshot
from app.services.cache import TTLCache
from app.services.headers import parse_subscription_userinfo


_fetch_cache: TTLCache[FetchResult] = TTLCache()
_provider_cache: TTLCache[ProviderDocument] = TTLCache()
_snapshot_cache: TTLCache[SourceSnapshot] = TTLCache()


async def fetch_source(name: str, source: SourceConfig) -> FetchResult:
    settings = get_settings()
    ttl = source.cache_ttl_seconds or settings.cache_ttl_seconds
    cached = _fetch_cache.get(name)
    if cached is not None:
        return cached

    headers = {"User-Agent": settings.default_user_agent}
    headers.update(source.headers)

    async with httpx.AsyncClient(timeout=settings.request_timeout_seconds, follow_redirects=True) as client:
        response = await client.get(source.url, headers=headers)
        response.raise_for_status()
        result = FetchResult(text=response.text, headers=dict(response.headers))
        _fetch_cache.set(name, result, ttl)
        return result


async def build_provider_document(name: str, source: SourceConfig) -> ProviderDocument:
    settings = get_settings()
    ttl = source.cache_ttl_seconds or settings.cache_ttl_seconds
    cache_key = f"provider:{name}"
    cached = _provider_cache.get(cache_key)
    if cached is not None:
        return cached

    fetched = await fetch_source(name, source)

    if source.kind != "clash_yaml":
        raise ValueError(f"Unsupported source kind: {source.kind}")

    proxies = parse_clash_yaml_proxies(fetched.text)
    proxies = transform_proxies(proxies, source, settings.max_proxy_name_length)

    document = ProviderDocument(proxies=proxies)
    _provider_cache.set(cache_key, document, ttl)
    return document


async def build_source_snapshot(name: str, source: SourceConfig) -> SourceSnapshot:
    settings = get_settings()
    ttl = source.cache_ttl_seconds or settings.cache_ttl_seconds
    cache_key = f"snapshot:{name}"
    cached = _snapshot_cache.get(cache_key)
    if cached is not None:
        return cached

    fetched = await fetch_source(name, source)
    document = await build_provider_document(name, source)
    snapshot = SourceSnapshot(
        name=name,
        display_name=source.display_name or name,
        document=document,
        headers=fetched.headers,
        quota=parse_subscription_userinfo(fetched.headers),
    )
    _snapshot_cache.set(cache_key, snapshot, ttl)
    return snapshot


async def build_source_snapshots(source_items: Iterable[tuple[str, SourceConfig]]) -> list[SourceSnapshot]:
    snapshots: list[SourceSnapshot] = []
    for name, source in source_items:
        snapshots.append(await build_source_snapshot(name, source))
    return snapshots


async def build_merged_provider_document(source_items: Iterable[tuple[str, SourceConfig]]) -> ProviderDocument:
    snapshots = await build_source_snapshots(source_items)
    proxies: list[dict[str, Any]] = []
    seen: set[str] = set()

    for snapshot in snapshots:
        for proxy in snapshot.document.proxies:
            candidate = dict(proxy)
            name = str(candidate.get("name", "")).strip()
            if not name:
                continue
            original = name
            index = 2
            while name in seen:
                name = f"{original} #{index}"
                index += 1
            candidate["name"] = name
            seen.add(name)
            proxies.append(candidate)

    return ProviderDocument(proxies=proxies)


async def get_first_quota(source_items: Iterable[tuple[str, SourceConfig]]):
    source_list = list(source_items)
    if not source_list:
        return None
    snapshot = await build_source_snapshot(source_list[0][0], source_list[0][1])
    return snapshot.quota


def parse_clash_yaml_proxies(text: str) -> list[dict[str, Any]]:
    data = yaml.safe_load(text)
    if not isinstance(data, dict):
        raise ValueError("Upstream YAML must be a mapping with a top-level 'proxies' field")

    proxies = data.get("proxies")
    if not isinstance(proxies, list):
        raise ValueError("Upstream YAML must contain a list field named 'proxies'")

    normalized: list[dict[str, Any]] = []
    for item in proxies:
        if not isinstance(item, dict):
            continue
        if not item.get("name") or not item.get("type"):
            continue
        normalized.append(item)
    return normalized


def transform_proxies(
    proxies: list[dict[str, Any]], source: SourceConfig, max_proxy_name_length: int
) -> list[dict[str, Any]]:
    include = re.compile(source.include_regex) if source.include_regex else None
    exclude = re.compile(source.exclude_regex) if source.exclude_regex else None

    transformed: list[dict[str, Any]] = []
    seen: dict[str, int] = {}

    for proxy in proxies:
        name = str(proxy.get("name", "")).strip()
        if not name:
            continue
        if include and not include.search(name):
            continue
        if exclude and exclude.search(name):
            continue

        new_proxy = dict(proxy)
        new_name = f"{source.prefix}{name}{source.suffix}".strip()
        if len(new_name) > max_proxy_name_length:
            new_name = new_name[:max_proxy_name_length].rstrip()

        count = seen.get(new_name, 0) + 1
        seen[new_name] = count
        if count > 1:
            new_name = f"{new_name} #{count}"

        new_proxy["name"] = new_name
        transformed.append(new_proxy)

    return transformed


def dump_provider_yaml(document: ProviderDocument) -> str:
    return yaml.safe_dump(
        {"proxies": document.proxies},
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=False,
    )