from typing import List, Tuple
import json
from threading import Semaphore
from alayaflow.component.llm_node import LLMComponent, ResponseFormat
from ..schemas import TaskState


def create_rewrite_query_node():
    MODEL_ID: str = "deepseek-chat"
    MAX_CONCURRENCY: int = 10
    limiter = Semaphore(MAX_CONCURRENCY)

    def _build_system_prompt() -> str:
        return """
        你是一个“混合检索（Hybrid Search）查询优化专家”。
        你的目标是生成一个**既能被向量模型理解语义，又能被关键词引擎（BM25）精准匹配**的查询。

        ### 核心策略
        1. **实体锚点（必须）**：必须从 `memory_json` 中提取核心实体（人名、公司名、项目名）。
        2. **关键词增强**：除了自然语言描述，**必须包含具体的领域术语**。
        3. **同义词扩展**：如果字段名有常用别名，请在句子中自然提及，以增加关键词命中率。
        - *示例*：查 `cv` -> Query 应包含 "简历" 和 "CV"。
        - *示例*：查 `revenue` -> Query 应包含 "营收" 和 "收入"。
        4. **高密度表达**：保持句子通顺，但**去除无意义的客套话**（如“请帮我找”、“关于...的文档”）。
        5. **幻觉控制**：绝对不要编造 memory 中不存在的具体数值。

        ### 输出格式
        严格 JSON: `{"query": "..."}`

        ### Few-Shot Examples (混合检索优化版)
        Context: Path="candidate/edu", Keys=["school", "gpa"], Memory={"name": "王强"}
        Output: {"query": "王强的毕业院校、大学成绩单及GPA绩点"}
        (分析：加入了"大学"、"成绩单"、"绩点"等高频同义词，利于关键词匹配)

        Context: Path="project/summary", Keys=["tech_stack"], Memory={"project": "AlphaGo"}
        Output: {"query": "AlphaGo项目的技术栈、使用的编程语言及框架架构"}
        (分析：扩展了"编程语言"、"框架"等具体维度的词)

        Context: Path="basic_info", Keys=["phone", "email"], Memory={"name": "Sarah"}
        Output: {"query": "Sarah的联系方式、手机号码及电子邮箱地址"}
        (分析：同时覆盖"联系方式"、"手机"、"邮箱"等潜在关键词)
        """.strip()

    def _build_user_prompt(
        original_query: str,
        path: Tuple[str, ...],
        keys: List[str],
        memory_json: str,
    ) -> str:
        path_context = " > ".join(path) if path else "根节点"
        keys_context = ", ".join(keys)

        return f"""
        ### 任务输入
        [待检索字段]: {keys_context}
        [字段归属路径]: {path_context}
        [上下文记忆 (Memory)]: {memory_json if memory_json else "暂无"}
        [原始输入]: {original_query}

        ### 指令
        请生成一个适合混合检索的 Query（自然语言与精准关键词并重）：
        """.strip()
        
    def rewrite_query_node(state: TaskState):
        task = state["task"]
        path = task.path
        keys = list(task.keys)
        
        query_parts = list(path) + keys
        original_query = "；".join([p for p in query_parts if p])

        memory_json = state.get("memory_json", "")

        try:
            with limiter:
                llm = LLMComponent(
                    model_id=MODEL_ID,
                    system_prompt=_build_system_prompt(),
                    prompt=_build_user_prompt(
                        original_query=original_query,
                        path=path,
                        keys=keys,
                        memory_json=memory_json,
                    ),
                    response_format=ResponseFormat.JSON,
                    temperature=0.0,
                )

                msg = llm()
                obj = json.loads(msg.content)

                rewritten = obj.get("query", "")
                if not isinstance(rewritten, str) or not rewritten.strip():
                    rewritten = original_query

                return {
                    "query": rewritten.strip(),
                }

        except Exception as e:
            old_errs = state.get("errors", {}) if isinstance(state.get("errors", {}), dict) else {}
            return {
                "query": original_query,
                "errors": {
                    **old_errs,
                    "rewrite_query": f"LLM Error: {type(e).__name__}: {e}",
                },
            }
    return rewrite_query_node