import json
from threading import Semaphore

from langgraph.runtime import Runtime

from alayaflow.component.llm_node import LLMComponent, ResponseFormat

from ..schemas import TaskState, WorkflowContext
from ..utils import make_patch


def create_generate_node():
    MAX_CONCURRENCY: int = 10
    limiter = Semaphore(MAX_CONCURRENCY)

    def _build_system_prompt(keys: list[str]) -> str:
        keys_str = ", ".join(keys)

        return f"""
你是一个严谨的“局部字段抽取器”（table patch extractor）。

你的任务是：**只为指定字段抽取值**，严格依据提供的记忆信息和知识片段，不得猜测或编造。

通用规则：
1. 输出必须是严格合法 JSON，不允许包含解释、Markdown、代码块或多余文本。
2. **只允许输出以下字段（不多不少）**：{keys_str}
3. 所有字段值必须是字符串。
4. 找不到 / 不确定 / 空值 / 占位符 → 必须输出空字符串 ""。
5. 字段名可能存在空格或轻微变体（如“姓 名”≈“姓名”），允许智能匹配，但不得扩展到未指定字段。

占位符识别：
- 若候选值是模板占位符或签字日期类文本
（如“签字： 年 月 日”“学院盖章： 年 月 日”等），必须返回 ""。""".strip()

    def _build_user_prompt(memory_json: str, content_text: str, path: list[str], keys: list[str],) -> str:
        path_str = " / ".join(path) if path else "<root>"
        keys_str = ", ".join(keys)

        json_skeleton = "{\n" + ",\n".join([f'  "{k}": ""' for k in keys]) + "\n}"

        return f"""
【本次任务定位】
字段路径（仅用于语义定位，不要输出）：{path_str}
需要抽取的字段：{keys_str}

【用户记忆】
{memory_json}

【知识库片段】
{content_text}

【输出要求】
- 只输出一个 JSON 对象
- key 必须严格为：{keys_str}
- 无法确定 / 空值 / 占位符 → 输出 ""

【JSON 输出模板】
{json_skeleton}""".strip()

    def _non_empty_kv(extracted: dict[str, str]) -> dict[str, str]:
        return {k: v for k, v in extracted.items() if isinstance(v, str) and v.strip()}
    
    def generate_node(state: TaskState, runtime: Runtime[WorkflowContext]):
        task = state["task"]
        path = task.path
        keys = list(task.keys)

        memory_json = state.get("memory_json", "")
        docs = state.get("docs", [])
        task_id = f"{'/'.join(path) or '<root>'}:{','.join(keys)}"
        default_kv = {k: "" for k in keys}
        default_patch = make_patch(path, default_kv)
        if not memory_json and not docs:
            return {
                "final_result": default_patch,
                "context_by_task": {task_id: []},
            }
        formatted_context = "\n\n".join(
            [f"片段 {i+1}:\n{doc}" for i, doc in enumerate(docs)]
        )
        try:
            with limiter:
                system_prompt = _build_system_prompt(keys)
                user_prompt = _build_user_prompt(memory_json, formatted_context, path, keys)

                llm = LLMComponent(
                    model_id=runtime.context.extract_model_id,
                    system_prompt=system_prompt,
                    prompt=user_prompt,
                    response_format=ResponseFormat.JSON,
                    temperature=0.6,
                )

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

                extracted = {}
                for k in keys:
                    v = obj.get(k, "")
                    extracted[k] = (str(v).strip() if v is not None else "")
                patch = make_patch(path, extracted)
                mem_kv = _non_empty_kv(extracted)
                mem_patch = make_patch(path, mem_kv) if mem_kv else {}
                return {
                    "final_result": patch,
                    "memory_patch": mem_patch,
                    "context_by_task": {task_id: docs},
                }

        except Exception as e:
            return {
                "final_result": default_patch,
                "memory_patch": {}, 
                "context_by_task": {task_id: docs},
                "errors": {task_id: f"LLM Error: {type(e).__name__}: {e}"},
            }

    return generate_node
