src/services/chat.py

"""多輪對話服務 — 透過 ModelRegistry 取得 LLM（Sprint D：預設 Qwen3-VL-8B）

【這個模組在做什麼？】
  接收多輪對話歷史，生成自然語言回應。
  支援 prompt-based tool calling：模型可輸出 JSON 工具呼叫，
  由上層（hub-service）執行後將結果回傳繼續對話。

【模型來源】（Sprint D, 2026-04-27）
  ModelRegistry.get("chat") → 預設 qwen3-vl-8b-instruct（Apache 2.0，dense 8B
  vision-language，256K context）；CHAT_MODEL=gemma-4-e4b 可回退 Gemma 4 rollback
  hatch（Apple Silicon dev / Qwen3-VL prod issue / git-bisect）。其餘 9 個 capability
  仍綁 gemma-4-e4b（split 不 replace），不再與 chat 共享同一 model instance。

【為什麼 tokenizer.apply_chat_template 與兩個 model 都通？】
  Qwen3-VL 用 AutoProcessor（multimodal），但 processor.tokenizer 支援純文字
  apply_chat_template，介面與 Gemma 4 / Qwen2.5 text tokenizer 一致。本檔走純文字
  path（無 image input），model-agnostic。tool calling 走 prompt-based JSON parse，
  不依賴 transformers native tool calling，跨模型相容。
"""

from __future__ import annotations

import json
import logging
import re

from pydantic import BaseModel, ConfigDict, model_validator


logger = logging.getLogger(__name__)

# ── 輸入大小限制（防止惡意超大 payload 造成 OOM） ────────────────────────────────
MAX_PAGE_CONTEXT_BYTES = 4096  # page_context JSON 序列化上限（位元組）
MAX_MESSAGES_COUNT = 50  # messages 陣列最大筆數
MAX_MESSAGE_CONTENT_CHARS = 4000  # 單則訊息 content 字串上限（字元）

# ── 資料模型 ──────────────────────────────────────────────────────────────────


class ChatRequest(BaseModel):
    """多輪對話請求

    v17 P2 — extra="forbid": defends against caller field-name drift.
    Existing model_validator already enforces size limits.
    """

    model_config = ConfigDict(extra="forbid")

    messages: list[dict]  # [{role, content, toolCalls?, toolName?}]
    page_context: dict  # {page: str, data: dict}
    tool_definitions: list[dict] | None = None  # 可覆寫預設工具清單

    @model_validator(mode="before")
    @classmethod
    def validate_payload_size(cls, values: dict) -> dict:
        """驗證輸入大小，避免攻擊者送超大 payload 導致 tokenization 階段 OOM。"""

        # 1. page_context JSON 序列化不得超過 4096 bytes
        page_ctx = values.get("page_context")
        if page_ctx is not None:
            ctx_bytes = len(
                json.dumps(page_ctx, ensure_ascii=False, default=str).encode()
            )
            if ctx_bytes > MAX_PAGE_CONTEXT_BYTES:
                msg = (
                    f"page_context 序列化後為 {ctx_bytes} bytes，"
                    f"超過上限 {MAX_PAGE_CONTEXT_BYTES} bytes"
                )
                raise ValueError(msg)

        # 2. messages 陣列不得超過 50 筆
        msgs = values.get("messages")
        if msgs is not None:
            if not isinstance(msgs, list):
                raise ValueError("messages 必須為陣列")
            if len(msgs) > MAX_MESSAGES_COUNT:
                msg = f"messages 共 {len(msgs)} 筆，超過上限 {MAX_MESSAGES_COUNT} 筆"
                raise ValueError(msg)

            # 3. 每則訊息的 content 不得超過 4000 字元
            for i, m in enumerate(msgs):
                content = m.get("content", "") if isinstance(m, dict) else ""
                if (
                    isinstance(content, str)
                    and len(content) > MAX_MESSAGE_CONTENT_CHARS
                ):
                    msg = (
                        f"messages[{i}].content 長度 {len(content)}，"
                        f"超過上限 {MAX_MESSAGE_CONTENT_CHARS} 字元"
                    )
                    raise ValueError(msg)

        return values


class ChatResponse(BaseModel):
    """對話回應（可能包含工具呼叫）"""

    content: str
    tool_calls: list[dict] | None = None  # [{name, arguments}]
    # "local_model" = 真模型輸出；"fallback" = 模型載入失敗時的 canned「warming up」回覆。
    # 讓 client 能 disclose 並提供重試。
    source: str = "local_model"


# ── 預設工具定義（嵌入 system prompt） ────────────────────────────────────────

DEFAULT_TOOLS = [
    "search_items(query: string) — Search user's wardrobe by text",
    "get_item_details(item_id: string) — Get details of a specific item",
    "search_marketplace(query: string, max_price?: number) — Search marketplace listings",
    "compare_items(item_ids: string[]) — Compare multiple items side by side",
    "get_price_info(item_id: string) — Get market price and trend for an item",
    "add_to_wishlist(item_id: string) — Add an item to user's wishlist",
    "get_recommendations(occasion?: string) — Get outfit recommendations",
    "filter_feed(query: string) — Filter feed content by topic",
    "get_user_taste_profile() — Get the user's style DNA: top brands, price range, wardrobe stats, style keywords",
    "get_brand_affinity(limit?: number) — Get top fashion brands with counts and values",
    "answer_persona_question(question: string) — Answer questions about the user's taste and style preferences",
]

MAX_HISTORY = 10  # 最大對話歷史長度（Qwen3-VL 256K 雖夠，但 prompt 控成本 + 延遲）
MAX_NEW_TOKENS = 512  # 最大生成 token 數


# ── System prompt 組裝 ───────────────────────────────────────────────────────


def _build_system_prompt(
    page_context: dict,
    tool_definitions: list[dict] | None,
) -> str:
    """組裝 system prompt，包含工具定義和頁面上下文。

    tool_definitions 語義：
      - None（未傳）→ 注入 DEFAULT_TOOLS（一般購物 copilot，行為不變）。
      - 非空 list → 用 caller 自訂工具清單。
      - **空 list `[]` → 完全不宣告工具**，連「respond with JSON block」指令都拿掉。
        這給「純文字生成」用途（如議價回覆草稿）：注入內容無法誘導模型吐 tool_calls
        JSON 進輸出，因為 system prompt 根本沒提過工具（security review P1, #150）。
    """
    if tool_definitions is not None and len(tool_definitions) == 0:
        page_desc = json.dumps(page_context, ensure_ascii=False, default=str)
        return (
            "You are a helpful fashion assistant for Wardrobe OS.\n"
            f"Current page context: {page_desc}\n\n"
            "Respond in natural conversational text. Be concise and helpful."
        )

    if tool_definitions:
        tool_lines = "\n".join(
            f"- {t.get('name', '?')}({t.get('params', '')}) — {t.get('description', '')}"
            for t in tool_definitions
        )
    else:
        tool_lines = "\n".join(f"- {t}" for t in DEFAULT_TOOLS)

    page_desc = json.dumps(page_context, ensure_ascii=False, default=str)

    return (
        "You are a helpful fashion assistant for Wardrobe OS.\n"
        "You can use tools to help answer questions. "
        "When you need to use a tool, respond with ONLY a JSON block:\n"
        '{"tool_calls": [{"name": "tool_name", "arguments": {"arg": "value"}}]}\n\n'
        f"Available tools:\n{tool_lines}\n\n"
        f"Current page context: {page_desc}\n\n"
        "When you have enough information, respond in natural conversational text. "
        "Be concise and helpful."
    )


# ── 對話歷史格式化 ───────────────────────────────────────────────────────────


def _format_messages(messages: list[dict]) -> list[dict]:
    """將前端訊息陣列轉換為 Qwen chat 格式。tool_result 包裝成 user 訊息。"""
    formatted: list[dict] = []
    for msg in messages[-MAX_HISTORY:]:
        role = msg.get("role", "user")
        content = msg.get("content", "")

        if role == "tool_result":
            # 工具結果包裝成 user 訊息，讓模型知道工具回傳了什麼
            tool_name = msg.get("toolName", "unknown")
            formatted.append(
                {
                    "role": "user",
                    "content": f"Tool '{tool_name}' returned: {content}",
                }
            )
        elif role in ("user", "assistant"):
            formatted.append({"role": role, "content": content})

    return formatted


# ── 工具呼叫解析 ─────────────────────────────────────────────────────────────


def _parse_tool_calls(text: str) -> list[dict] | None:
    """
    嘗試從模型輸出解析 tool_calls JSON。
    策略：直接 JSON 解析 → regex 抽取 → None（純文字回應）
    """
    text = text.strip()

    # 剝掉 markdown code fence
    if "```" in text:
        lines = text.split("\n")
        text = "\n".join(ln for ln in lines if not ln.strip().startswith("```"))
        text = text.strip()

    # 策略一：直接解析整段文字
    try:
        data = json.loads(text)
        if isinstance(data, dict) and "tool_calls" in data:
            calls = data["tool_calls"]
            if isinstance(calls, list) and len(calls) > 0:
                return calls
    except (json.JSONDecodeError, ValueError):
        pass

    # 策略二：regex 抽取 {"tool_calls": [...]} 區塊
    match = re.search(r'\{"tool_calls"\s*:\s*\[.*?\]\s*\}', text, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group())
            calls = data.get("tool_calls", [])
            if isinstance(calls, list) and len(calls) > 0:
                return calls
        except (json.JSONDecodeError, ValueError):
            pass

    return None


# ── 模型存取（透過 ModelRegistry 集中管理） ──────────────────────────────────


async def _get_chat_model():
    """從 Registry 取得 capability `chat` 模型實例（Sprint D：預設 Qwen3-VL-8B）。"""
    from src.registry import get_registry

    reg = get_registry()
    loaded = await reg.get("chat")
    return loaded.model, loaded.tokenizer


# ── 主函式 ───────────────────────────────────────────────────────────────────


async def chat_respond(request: ChatRequest) -> ChatResponse:
    """
    多輪對話回應（含 tool calling 支援）。

    流程：組裝 system prompt → 從 Registry 取得模型 → 推論 → 解析 tool_calls 或回純文字
    """
    import asyncio

    # 空訊息檢查
    if not request.messages:
        return ChatResponse(content="Please send a message to start the conversation.")

    # 從 Registry 取得模型
    try:
        model, tokenizer = await _get_chat_model()
    except Exception as e:
        logger.warning("Chat 模型不可用: %s", e)
        return ChatResponse(
            content="I'm still warming up. Please try again in a moment.",
            source="fallback",
        )

    # 組裝對話
    system_prompt = _build_system_prompt(request.page_context, request.tool_definitions)
    chat_messages = [{"role": "system", "content": system_prompt}]
    chat_messages.extend(_format_messages(request.messages))

    def _inference():
        import torch

        # Tokenize + 推論
        text = tokenizer.apply_chat_template(
            chat_messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = tokenizer(text, return_tensors="pt").to(model.device)

        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
            )

        # 只取新生成的 token
        generated = output[0][inputs["input_ids"].shape[1] :]
        return tokenizer.decode(generated, skip_special_tokens=True).strip()

    result_text = await asyncio.to_thread(_inference)

    # 嘗試解析工具呼叫
    tool_calls = _parse_tool_calls(result_text)
    if tool_calls:
        return ChatResponse(content="", tool_calls=tool_calls)

    return ChatResponse(content=result_text)