ForcePilot/backend/package/yuxi/channel/sdk/targets/pipeline.py

import logging
import time
from collections.abc import Callable

from yuxi.channel.protocols import DirectoryEntry
from yuxi.channel.sdk.targets.normalizer import (
    detect_target_kind,
    looks_like_target_id,
    normalize_target_input,
)
from yuxi.channel.sdk.targets.types import (
    AmbiguousMode,
    ResolveError,
    ResolveErrorKind,
    ResolvedTarget,
)

logger = logging.getLogger(__name__)

PluginTargetResolver = Callable[[str, str], ResolvedTarget | None]


class TargetResolverPipeline:
    def __init__(
        self,
        *,
        directory=None,
        directory_config: dict | None = None,
        directory_account_id: str | None = None,
        plugin_resolver: PluginTargetResolver | None = None,
        plugin_hint: str = "",
        ambiguous_mode: AmbiguousMode = AmbiguousMode.ERROR,
        cache_ttl: float = 1800.0,
    ):
        self._directory = directory
        self._directory_config = directory_config or {}
        self._directory_account_id = directory_account_id
        self._plugin_resolver = plugin_resolver
        self._plugin_hint = plugin_hint
        self._ambiguous_mode = ambiguous_mode
        self._cache: dict[str, tuple[float, list[DirectoryEntry]]] = {}
        self._cache_ttl = cache_ttl

    async def resolve(self, raw_input: str) -> ResolvedTarget:
        normalized = normalize_target_input(raw_input)
        if not normalized:
            raise ResolveError(
                kind=ResolveErrorKind.EMPTY_INPUT,
                message="目标输入为空",
                hint=self._plugin_hint,
            )

        kind = detect_target_kind(normalized)

        if looks_like_target_id(normalized):
            return ResolvedTarget(
                to=normalized,
                kind=kind,
                display=normalized,
                source="normalized",
            )

        if self._directory is not None:
            try:
                entry = await self._search_directory(normalized, kind)
                if entry is not None:
                    return entry
            except ResolveError as e:
                if e.kind == ResolveErrorKind.AMBIGUOUS and self._ambiguous_mode != AmbiguousMode.ERROR:
                    best = self._pick_ambiguous(e.candidates, self._ambiguous_mode)
                    if best is not None:
                        return ResolvedTarget(
                            to=best.id,
                            kind=kind,
                            display=best.name or best.handle or "",
                            source="directory",
                            confidence=0.5,
                        )
                raise

        if self._plugin_resolver is not None:
            fallback = self._plugin_resolver(normalized, kind)
            if fallback is not None:
                return ResolvedTarget(
                    to=fallback.to,
                    kind=kind,
                    display=fallback.display or normalized,
                    source="fallback",
                    confidence=0.3,
                )

        raise ResolveError(
            kind=ResolveErrorKind.NOT_FOUND,
            message=f"未找到目标: {normalized}",
            hint=self._plugin_hint,
        )

    async def _search_directory(self, query: str, kind: str) -> ResolvedTarget | None:
        cache_key = f"{kind}"
        now = time.monotonic()

        if cache_key in self._cache:
            ts, entries = self._cache[cache_key]
            if now - ts < self._cache_ttl:
                return self._match_entries(query, entries, kind)

        try:
            if kind == "user":
                entries = await self._directory.list_peers_live(
                    config=self._directory_config,
                    account_id=self._directory_account_id,
                    query=query,
                    limit=20,
                )
            else:
                entries = await self._directory.list_groups_live(
                    config=self._directory_config,
                    account_id=self._directory_account_id,
                    query=query,
                    limit=20,
                )
        except Exception:
            logger.debug("Directory lookup failed for query=%s kind=%s", query, kind)
            return None

        self._cache[cache_key] = (now, entries)
        return self._match_entries(query, entries, kind)

    def _match_entries(self, query: str, entries: list[DirectoryEntry], kind: str) -> ResolvedTarget | None:
        matches = [entry for entry in entries if self._entry_matches(query, entry)]

        if len(matches) == 1:
            entry = matches[0]
            return ResolvedTarget(
                to=entry.id,
                kind=kind,
                display=entry.name or entry.handle or "",
                source="directory",
            )

        if len(matches) > 1:
            if self._ambiguous_mode == AmbiguousMode.ERROR:
                raise ResolveError(
                    kind=ResolveErrorKind.AMBIGUOUS,
                    message=f"找到 {len(matches)} 个匹配目标",
                    candidates=matches,
                    hint=self._plugin_hint,
                )
            best = self._pick_ambiguous(matches, self._ambiguous_mode)
            if best is not None:
                return ResolvedTarget(
                    to=best.id,
                    kind=kind,
                    display=best.name or best.handle or "",
                    source="directory",
                    confidence=0.5,
                )

        return None

    @staticmethod
    def _entry_matches(query: str, entry: DirectoryEntry) -> bool:
        q = query.lower().strip().lstrip("@").lstrip("#")
        candidates = [
            entry.id.lower(),
            (entry.name or "").lower(),
            (entry.handle or "").lower(),
        ]
        return any(q in c for c in candidates if c)

    @staticmethod
    def _pick_ambiguous(candidates: list[DirectoryEntry], mode: AmbiguousMode) -> DirectoryEntry | None:
        if not candidates:
            return None
        if mode == AmbiguousMode.FIRST:
            return candidates[0]
        ranked = sorted(
            candidates,
            key=lambda e: e.rank if isinstance(e.rank, (int, float)) else 0,
            reverse=True,
        )
        return ranked[0]