ForcePilot/backend/package/yuxi/channels/adapters/wechat/selector.py

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any, Literal


@dataclass
class SelectorOption:
    id: str
    text: str
    is_default: bool = False


@dataclass
class SelectorItem:
    question_key: str
    title: str = ""
    options: list[SelectorOption] | None = None
    type: Literal["single", "multi"] = "single"
    disable: bool = False


@dataclass
class SelectorCard:
    task_id: str
    title: str = "请选择"
    description: str = ""
    items: list[SelectorItem] | None = None
    created_at: float = 0.0
    expires_at: float = 0.0


class WeChatSelectorAdapter:
    def __init__(self, default_ttl_hours: int = 24):
        self._default_ttl_hours = default_ttl_hours
        self._active_selectors: dict[str, SelectorCard] = {}

    @staticmethod
    def build_selector(
        task_id: str,
        title: str,
        items: list[SelectorItem],
        description: str = "",
        ttl_hours: int = 24,
    ) -> SelectorCard:
        now = time.time()
        return SelectorCard(
            task_id=task_id,
            title=title,
            description=description,
            items=items,
            created_at=now,
            expires_at=now + ttl_hours * 3600,
        )

    @staticmethod
    def build_instance_option(id: str, text: str, is_default: bool = False) -> SelectorOption:
        return SelectorOption(id=id, text=text, is_default=is_default)

    @staticmethod
    def build_instance_item(
        question_key: str,
        title: str,
        options: list[SelectorOption],
        item_type: Literal["single", "multi"] = "single",
    ) -> SelectorItem:
        return SelectorItem(
            question_key=question_key,
            title=title,
            options=options,
            type=item_type,
        )

    def create_selector(self, card: SelectorCard) -> str:
        self._active_selectors[card.task_id] = card
        return card.task_id

    def get_selector(self, task_id: str) -> SelectorCard | None:
        card = self._active_selectors.get(task_id)
        if card and card.expires_at < time.time():
            self._active_selectors.pop(task_id, None)
            return None
        return card

    def is_expired(self, task_id: str) -> bool:
        card = self._active_selectors.get(task_id)
        if card is None:
            return True
        if card.expires_at < time.time():
            self._active_selectors.pop(task_id, None)
            return True
        return False

    def resolve_selections(self, task_id: str, selections: dict[str, list[str]]) -> dict[str, Any] | None:
        card = self.get_selector(task_id)
        if card is None or not card.items:
            return None

        resolved: dict[str, Any] = {"task_id": task_id, "answers": {}}
        for item in card.items:
            key = item.question_key
            if key not in selections:
                continue
            picked = selections[key]
            if item.options:
                valid_ids = {opt.id for opt in item.options}
                filtered = [v for v in picked if v in valid_ids]
                resolved["answers"][key] = filtered if item.type == "multi" else (filtered[0] if filtered else "")

        return resolved

    def render_taskcard_payload(self, to_user: str, agent_id: str, card: SelectorCard) -> dict[str, Any]:
        select_list: list[dict[str, Any]] = []
        if card.items:
            for item in card.items:
                option_list: list[dict[str, Any]] = []
                if item.options:
                    option_list = [
                        {"id": opt.id, "text": opt.text, "is_default": opt.is_default} for opt in item.options
                    ]
                select_list.append(
                    {
                        "question_key": item.question_key,
                        "title": item.title,
                        "type": item.type,
                        "disable": item.disable,
                        "option_list": option_list,
                    }
                )

        return {
            "touser": to_user,
            "agentid": agent_id,
            "msgtype": "template_card",
            "template_card": {
                "card_type": "button_interaction",
                "main_title": {"title": card.title, "desc": card.description},
                "task_id": card.task_id,
                "select_list": select_list,
                "submit_button": {
                    "text": "确认",
                    "key": f"submit_{card.task_id}",
                },
            },
        }

    def list_active_selectors(self) -> list[dict[str, Any]]:
        now = time.time()
        expired = [tid for tid, c in self._active_selectors.items() if c.expires_at < now]
        for tid in expired:
            self._active_selectors.pop(tid, None)

        return [
            {
                "task_id": c.task_id,
                "title": c.title,
                "expires_in": max(0, int(c.expires_at - now)),
            }
            for c in self._active_selectors.values()
        ]