core/plugin/base.py

"""
插件抽象与注册表：type_name -> 插件实现。

三层设计：
  AbstractPlugin   — 最底层接口，理论上支持任意协议（非 Cookie、非 SSE 的站点也能接）。
  BaseSitePlugin   — Cookie 认证 + SSE 流式站点的通用编排，插件开发者继承它只需实现 5 个 hook。
  PluginRegistry   — 全局注册表。
"""

import json
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, AsyncIterator

from playwright.async_api import BrowserContext, Page

from core.api.schemas import InputAttachment
from core.config.settings import get
from core.plugin.errors import (  # noqa: F401  — re-export for backward compat
    AccountFrozenError,
    BrowserResourceInvalidError,
)
from core.plugin.helpers import (
    apply_cookie_auth,
    create_page_for_site,
    stream_completion_via_sse,
)

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ResolvedModel:
    public_model: str
    upstream_model: str


# ---------------------------------------------------------------------------
# SiteConfig：纯声明式站点配置
# ---------------------------------------------------------------------------


@dataclass
class SiteConfig:
    """Cookie 认证站点的声明式配置，插件开发者只需填字段，无需写任何方法。"""

    start_url: str
    api_base: str
    cookie_name: str
    cookie_domain: str
    auth_keys: list[str]
    config_section: str = (
        ""  # config.yaml 中的 section，如 "claude"，用于覆盖 start_url/api_base
    )


# ---------------------------------------------------------------------------
# AbstractPlugin — 最底层抽象接口
# ---------------------------------------------------------------------------


class AbstractPlugin(ABC):
    """
    各 type（如 claude、kimi）需实现此接口并注册。
    若站点基于 Cookie + SSE，推荐直接继承 BaseSitePlugin 而非此类。
    """

    def __init__(self) -> None:
        self._session_state: dict[str, dict[str, Any]] = {}

    type_name: str

    async def create_page(
        self, context: BrowserContext, reuse_page: Page | None = None
    ) -> Page:
        raise NotImplementedError

    async def apply_auth(
        self,
        context: BrowserContext,
        page: Page,
        auth: dict[str, Any],
        *,
        reload: bool = True,
        **kwargs: Any,
    ) -> None:
        raise NotImplementedError

    async def ensure_request_ready(
        self,
        context: BrowserContext,
        page: Page,
        *,
        request_id: str = "",
        session_id: str | None = None,
        phase: str = "",
        account_id: str = "",
    ) -> None:
        del context, page, request_id, session_id, phase, account_id

    async def create_conversation(
        self,
        context: BrowserContext,
        page: Page,
        **kwargs: Any,
    ) -> str | None:
        raise NotImplementedError

    async def stream_completion(
        self,
        context: BrowserContext,
        page: Page,
        session_id: str,
        message: str,
        **kwargs: Any,
    ) -> AsyncIterator[str]:
        if False:
            yield  # 使抽象方法为 async generator，与子类一致，便于 async for 迭代
        raise NotImplementedError

    def parse_session_id(self, messages: list[dict[str, Any]]) -> str | None:
        return None

    def is_stream_end_event(self, payload: str) -> bool:
        """判断某条流式 payload 是否表示本轮响应已正常结束。默认不识别。"""
        return False

    def has_session(self, session_id: str) -> bool:
        return session_id in self._session_state

    def drop_session(self, session_id: str) -> None:
        self._session_state.pop(session_id, None)

    def drop_sessions(self, session_ids: list[str] | set[str]) -> None:
        for session_id in session_ids:
            self._session_state.pop(session_id, None)

    def model_mapping(self) -> dict[str, str] | None:
        """子类可覆盖；BaseSitePlugin 会从 config_section 的 model_mapping 读取。"""
        return None

    def normalized_model_mapping(self) -> dict[str, str]:
        mapping = self.model_mapping()
        if not isinstance(mapping, dict) or not mapping:
            raise ValueError("model_mapping is not implemented")
        normalized: dict[str, str] = {}
        for public_model, upstream_model in mapping.items():
            public_id = str(public_model or "").strip()
            upstream_id = str(upstream_model or "").strip()
            if public_id and upstream_id:
                normalized[public_id] = upstream_id
        if not normalized:
            raise ValueError("model_mapping is not implemented")
        return normalized

    def listed_model_mapping(self) -> dict[str, str]:
        return self.normalized_model_mapping()

    def default_public_model(self) -> str:
        listed = self.listed_model_mapping()
        if listed:
            return next(iter(listed))
        return next(iter(self.normalized_model_mapping()))

    def resolve_model(self, model: str | None) -> ResolvedModel:
        mapping = self.normalized_model_mapping()
        requested = str(model or "").strip()
        if not requested:
            default_public = self.default_public_model()
            return ResolvedModel(
                public_model=default_public,
                upstream_model=mapping[default_public],
            )
        if requested in mapping:
            return ResolvedModel(
                public_model=requested,
                upstream_model=mapping[requested],
            )
        for public_model, upstream_model in mapping.items():
            if requested == upstream_model:
                return ResolvedModel(
                    public_model=public_model,
                    upstream_model=upstream_model,
                )
        supported = ", ".join(mapping.keys())
        raise ValueError(f"Unknown model: {requested}. Supported models: {supported}")

    def on_http_error(self, message: str, headers: dict[str, str] | None) -> int | None:
        return None

    def stream_transport(self) -> str:
        return "page_fetch"

    def stream_transport_options(
        self,
        context: BrowserContext,
        page: Page,
        session_id: str,
        state: dict[str, Any],
        **kwargs: Any,
    ) -> dict[str, Any]:
        del context, page, session_id, state
        options: dict[str, Any] = {}
        proxy_url = str(kwargs.get("proxy_url") or "").strip()
        proxy_auth = kwargs.get("proxy_auth")
        if proxy_url:
            options["proxy_url"] = proxy_url
        if isinstance(proxy_auth, tuple) and len(proxy_auth) == 2:
            options["proxy_auth"] = proxy_auth
        return options


# ---------------------------------------------------------------------------
# BaseSitePlugin — Cookie + SSE 站点的通用编排
# ---------------------------------------------------------------------------


class BaseSitePlugin(AbstractPlugin):
    """
    Cookie 认证 + SSE 流式站点的公共基类。

    插件开发者继承此类后，只需：
      1. 声明 site = SiteConfig(...)        — 站点配置
      2. 实现 fetch_site_context()           — 获取站点上下文（如 org/user 信息）
      3. 实现 create_session()              — 调用站点 API 创建会话
      4. 实现 build_completion_url/body()    — 拼补全请求的 URL 与 body
      5. 实现 parse_stream_event()          — 解析单条流式事件（如 SSE data）

    create_page / apply_auth / create_conversation / stream_completion
    均由基类自动编排，无需重写。
    """

    site: SiteConfig  # 子类必须赋值

    # ---- 从 config.yaml 读取的 URL 属性（config_section 有值时覆盖默认） ----

    @property
    def start_url(self) -> str:
        if self.site.config_section:
            url = get(self.site.config_section, "start_url")
            if url:
                return str(url).strip()
        return self.site.start_url

    @property
    def api_base(self) -> str:
        if self.site.config_section:
            base = get(self.site.config_section, "api_base")
            if base:
                return str(base).strip()
        return self.site.api_base

    def model_mapping(self) -> dict[str, str] | None:
        """从 config 的 config_section.model_mapping 读取；未配置时返回 None。"""
        if self.site.config_section:
            m = get(self.site.config_section, "model_mapping")
            if isinstance(m, dict) and m:
                return {str(k): str(v) for k, v in m.items()}
        return None

    # ---- 基类全自动实现，子类无需碰 ----

    async def create_page(
        self,
        context: BrowserContext,
        reuse_page: Page | None = None,
    ) -> Page:
        return await create_page_for_site(
            context, self.start_url, reuse_page=reuse_page
        )

    async def apply_auth(
        self,
        context: BrowserContext,
        page: Page,
        auth: dict[str, Any],
        *,
        reload: bool = True,
    ) -> None:
        await apply_cookie_auth(
            context,
            page,
            auth,
            self.site.cookie_name,
            self.site.auth_keys,
            self.site.cookie_domain,
            reload=reload,
        )

    async def create_conversation(
        self,
        context: BrowserContext,
        page: Page,
        **kwargs: Any,
    ) -> str | None:
        extra_kwargs = dict(kwargs)
        request_id = str(extra_kwargs.pop("request_id", "") or "")
        # 调用子类获取站点上下文
        site_context = await self.fetch_site_context(
            context,
            page,
            request_id=request_id,
        )
        if site_context is None:
            logger.warning(
                "[%s] fetch_site_context 返回 None，请确认已登录", self.type_name
            )
            return None
        # 通过站点上下文创建会话
        conv_id = await self.create_session(
            context,
            page,
            site_context,
            request_id=request_id,
            **extra_kwargs,
        )
        if conv_id is None:
            return None
        state: dict[str, Any] = {"site_context": site_context}
        if kwargs.get("timezone") is not None:
            state["timezone"] = kwargs["timezone"]
        public_model = str(kwargs.get("public_model") or "").strip()
        if public_model:
            state["public_model"] = public_model
        upstream_model = str(kwargs.get("upstream_model") or "").strip()
        if upstream_model:
            state["upstream_model"] = upstream_model
        self._session_state[conv_id] = state
        logger.info(
            "[%s] create_conversation done conv_id=%s sessions=%s",
            self.type_name,
            conv_id,
            list(self._session_state.keys()),
        )
        return conv_id

    async def stream_completion(
        self,
        context: BrowserContext,
        page: Page,
        session_id: str,
        message: str,
        **kwargs: Any,
    ) -> AsyncIterator[str]:
        state = self._session_state.get(session_id)
        if not state:
            raise RuntimeError(f"未知会话 ID: {session_id}")

        request_id: str = kwargs.get("request_id", "")
        url = self.build_completion_url(session_id, state)
        attachments = list(kwargs.get("attachments") or [])
        prepared_attachments = await self.prepare_attachments(
            context,
            page,
            session_id,
            state,
            attachments,
            request_id=request_id,
        )
        body = self.build_completion_body(
            message,
            session_id,
            state,
            prepared_attachments,
        )
        body_json = json.dumps(body)

        logger.info(
            "[%s] stream_completion session_id=%s url=%s",
            self.type_name,
            session_id,
            url,
        )

        out_message_ids: list[str] = []
        transport_options = self.stream_transport_options(
            context,
            page,
            session_id,
            state,
            request_id=request_id,
            attachments=attachments,
            proxy_url=kwargs.get("proxy_url"),
            proxy_auth=kwargs.get("proxy_auth"),
        )

        async for text in stream_completion_via_sse(
            context,
            page,
            url,
            body_json,
            self.parse_stream_event,
            request_id,
            on_http_error=self.on_http_error,
            is_terminal_event=self.is_stream_end_event,
            collect_message_id=out_message_ids,
            transport=self.stream_transport(),
            transport_options=transport_options,
        ):
            yield text

        if out_message_ids and session_id in self._session_state:
            self.on_stream_completion_finished(session_id, out_message_ids)

    # ---- 子类必须实现的 hook ----

    @abstractmethod
    async def fetch_site_context(
        self,
        context: BrowserContext,
        page: Page,
        request_id: str = "",
    ) -> dict[str, Any] | None:
        """获取站点上下文信息（如 org_uuid、user_id 等），失败返回 None。"""
        del request_id
        ...

    @abstractmethod
    async def create_session(
        self,
        context: BrowserContext,
        page: Page,
        site_context: dict[str, Any],
        **kwargs: Any,
    ) -> str | None:
        """调用站点 API 创建会话，返回会话 ID，失败返回 None。"""
        ...

    @abstractmethod
    def build_completion_url(self, session_id: str, state: dict[str, Any]) -> str:
        """根据会话状态拼出补全请求的完整 URL。"""
        ...

    @abstractmethod
    def build_completion_body(
        self,
        message: str,
        session_id: str,
        state: dict[str, Any],
        prepared_attachments: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """构建补全请求体，返回 dict（基类负责 json.dumps）。"""
        ...

    @abstractmethod
    def parse_stream_event(
        self,
        payload: str,
    ) -> tuple[list[str], str | None, str | None]:
        """
        解析单条流式事件 payload（如 SSE data 行）。
        返回 (texts, message_id, error_message)。
        """
        ...

    # ---- 子类可选覆盖的 hook（有合理默认值） ----

    def on_stream_completion_finished(
        self,
        session_id: str,
        message_ids: list[str],
    ) -> None:
        """Hook：流式补全结束后调用，子类可按需用 message_ids 更新会话 state（如记续写用的父消息 id）。"""

    async def prepare_attachments(
        self,
        context: BrowserContext,
        page: Page,
        session_id: str,
        state: dict[str, Any],
        attachments: list[InputAttachment],
        request_id: str = "",
    ) -> dict[str, Any]:
        del context, page, session_id, state, attachments, request_id
        return {}


# ---------------------------------------------------------------------------
# PluginRegistry — 全局注册表
# ---------------------------------------------------------------------------


class PluginRegistry:
    """全局插件注册表：type_name -> AbstractPlugin。"""

    _plugins: dict[str, AbstractPlugin] = {}

    @classmethod
    def register(cls, plugin: AbstractPlugin) -> None:
        cls._plugins[plugin.type_name] = plugin

    @classmethod
    def get(cls, type_name: str) -> AbstractPlugin | None:
        return cls._plugins.get(type_name)

    @classmethod
    def resolve_model(cls, type_name: str, model: str | None) -> ResolvedModel:
        plugin = cls.get(type_name)
        if plugin is None:
            raise ValueError(f"Unknown provider: {type_name}")
        return plugin.resolve_model(model)

    @classmethod
    def model_metadata(cls, type_name: str) -> dict[str, Any]:
        plugin = cls.get(type_name)
        if plugin is None:
            raise ValueError(f"Unknown provider: {type_name}")
        mapping = plugin.listed_model_mapping()
        return {
            "provider": type_name,
            "public_models": list(mapping.keys()),
            "model_mapping": mapping,
            "default_model": plugin.default_public_model(),
        }

    @classmethod
    def all_types(cls) -> list[str]:
        return list(cls._plugins.keys())