"""QueryTransformer - RAG 查询改写

将用户原始查询改写为更适合知识库检索的形式：
- LLMQueryTransformer: 基于 LLM 的智能改写
- RuleQueryTransformer: 基于规则的改写（去停用词、同义扩展）
- NoOpQueryTransformer: 不改写，原样返回
"""

import json
import logging
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class TransformedQuery:
    """改写后的查询"""

    main_query: str
    sub_queries: list[str]
    original_query: str


class QueryTransformerBase(ABC):
    """查询改写抽象基类"""

    @abstractmethod
    async def transform(self, query: str) -> TransformedQuery:
        """改写查询"""
        ...


class LLMQueryTransformer(QueryTransformerBase):
    """基于 LLM 的查询改写

    通过 LLM 提取核心意图、分解子查询、添加领域术语。
    """

    def __init__(self, llm_gateway, max_sub_queries: int = 3):
        self._llm_gateway = llm_gateway
        self._max_sub_queries = max_sub_queries

    async def transform(self, query: str) -> TransformedQuery:
        """使用 LLM 改写查询"""
        prompt = (
            "You are a query rewriting assistant for a knowledge base retrieval system.\n"
            "Given a user query, your task is to:\n"
            "1. Extract the core intent of the query\n"
            "2. If the query is complex, decompose it into simpler sub-queries\n"
            "3. Add domain-specific terms that may improve retrieval\n\n"
            f"Original query: {query}\n\n"
            'Respond ONLY with a JSON object in this exact format: {"main_query": "...", "sub_queries": [...]}\n'
            "The main_query should be a concise, retrieval-optimized version of the original query.\n"
            "The sub_queries should be a list of simpler queries (0-3 items) that cover different aspects.\n"
            "Do not include any other text or explanation."
        )

        try:
            response = await self._llm_gateway.chat(
                messages=[{"role": "user", "content": prompt}],
                model="default",
            )
            data = json.loads(response.content)
            main_query = str(data.get("main_query", query))
            sub_queries = list(data.get("sub_queries", []))[: self._max_sub_queries]
            return TransformedQuery(
                main_query=main_query,
                sub_queries=sub_queries,
                original_query=query,
            )
        except Exception:
            logger.warning("LLM query transformation failed, falling back to original query")
            return TransformedQuery(
                main_query=query,
                sub_queries=[],
                original_query=query,
            )


class RuleQueryTransformer(QueryTransformerBase):
    """基于规则的查询改写

    去除填充词、提取关键名词短语、同义扩展。
    """

    _FILLER_WORDS_CN: list[str] = [
        "帮我", "请", "一下", "分析", "看看", "告诉我", "想知道", "请问",
    ]
    _FILLER_WORDS_EN: list[str] = [
        "please", "can you", "help me", "could you", "i want to", "i need to",
    ]

    def __init__(
        self,
        synonyms: dict[str, list[str]] | None = None,
        max_sub_queries: int = 3,
    ):
        self._synonyms = synonyms or {}
        self._max_sub_queries = max_sub_queries
        # Pre-compile filler patterns
        self._filler_patterns_cn = [
            re.compile(re.escape(w)) for w in self._FILLER_WORDS_CN
        ]
        self._filler_patterns_en = [
            re.compile(re.escape(w), re.IGNORECASE) for w in self._FILLER_WORDS_EN
        ]

    async def transform(self, query: str) -> TransformedQuery:
        """基于规则改写查询"""
        cleaned = query

        # Remove Chinese filler words
        for pattern in self._filler_patterns_cn:
            cleaned = pattern.sub("", cleaned)

        # Remove English filler words
        for pattern in self._filler_patterns_en:
            cleaned = pattern.sub("", cleaned)

        # Collapse whitespace
        cleaned = re.sub(r"\s+", " ", cleaned).strip()

        # If nothing left after cleaning, use original
        if not cleaned:
            cleaned = query

        # Synonym expansion
        sub_queries: list[str] = []
        for term, expansions in self._synonyms.items():
            if term in cleaned:
                for expansion in expansions:
                    if expansion != cleaned:
                        sub_queries.append(cleaned.replace(term, expansion))
                    if len(sub_queries) >= self._max_sub_queries:
                        break
            if len(sub_queries) >= self._max_sub_queries:
                break

        return TransformedQuery(
            main_query=cleaned,
            sub_queries=sub_queries,
            original_query=query,
        )


class NoOpQueryTransformer(QueryTransformerBase):
    """不做任何改写，原样返回"""

    async def transform(self, query: str) -> TransformedQuery:
        return TransformedQuery(
            main_query=query,
            sub_queries=[],
            original_query=query,
        )


def create_query_transformer(
    strategy: str = "none",
    llm_gateway=None,
    synonyms: dict[str, list[str]] | None = None,
    max_sub_queries: int = 3,
) -> QueryTransformerBase:
    """工厂函数：根据策略创建查询改写器"""
    if strategy == "llm":
        if llm_gateway is None:
            logger.warning("LLM strategy requested but no llm_gateway provided, falling back to NoOp")
            return NoOpQueryTransformer()
        return LLMQueryTransformer(llm_gateway, max_sub_queries=max_sub_queries)
    elif strategy == "rule":
        return RuleQueryTransformer(synonyms=synonyms, max_sub_queries=max_sub_queries)
    else:
        return NoOpQueryTransformer()
