import json
from typing import Any, Dict, List, Optional, Tuple
from collections import defaultdict
from threading import Semaphore

from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from langchain_core.runnables import RunnableConfig

from alayaflow.component.llm_node import LLMComponent, ResponseFormat
from alayaflow.clients.alayamem.http_client import HttpAlayaMemClient
from alayaflow.component.retrieve_node import RetrieveComponent

from .schemas import (
    FieldSpec, GroupTask, WorkflowInitArgs, 
    OverallState, TaskState, WorkflowContext, Input, Output
)
from .utils import make_patch, slim_docs


def flatten_leaf_tasks(specs: List[FieldSpec], base_path: Optional[List[str]] = None) -> List[Tuple[Tuple[str, ...], str]]:
    """
    返回：[(path_tuple, leaf_key), ...]
    """
    base_path = base_path or []
    out: List[Tuple[Tuple[str, ...], str]] = []

    def _as_list(x: Any) -> List[Any]:
        if x is None:
            return []
        if isinstance(x, list):
            return x
        return [x]

    for item in specs or []:
        if isinstance(item, str):
            out.append((tuple(base_path), item))
            continue

        if isinstance(item, dict):
            for parent, children in item.items():
                for child in _as_list(children):
                    if isinstance(child, str):
                        out.append((tuple(base_path + [parent]), child))
                    elif isinstance(child, dict):
                        out.extend(flatten_leaf_tasks([child], base_path + [parent]))
                    else:
                        pass
            continue

    return out


def create_plan_node():
    """创建规划节点：将输入的字段模板规划为多个独立的抽取任务"""
    def plan_node(state: OverallState, config: RunnableConfig):
        leaf = flatten_leaf_tasks(state["fields"])
        grouped: Dict[Tuple[str, ...], List[str]] = defaultdict(list)
        for path, key in leaf:
            grouped[path].append(key)

        tasks: List[GroupTask] = []
        for path, keys in grouped.items():
            # 去重保持顺序
            seen = set()
            uniq = []
            for k in keys:
                if k not in seen:
                    seen.add(k)
                    uniq.append(k)
            tasks.append(GroupTask(path=path, keys=tuple(uniq)))

        tasks.sort(key=lambda t: (len(t.path), t.path))
        return {"tasks": tasks}
    
    return plan_node


def create_extract_task_node(client: HttpAlayaMemClient):
    TOP_K: int = 5
    MAX_DOC_CHARS: int = 400
    MODEL_ID: str = "deepseek-chat"
    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. 字段名可能存在空格或轻微变体（如“姓 名”≈“姓名”），允许智能匹配，但不得扩展到未指定字段。
            
            长文本字段格式规则（必须遵守）：
            - 当字段内容包含**多个条目、多个时间段或多段经历**时：
            - 必须使用序号列表格式。
            - **每个条目占一行，条目之间必须使用 "\n" 换行符分隔。**
            - 不允许使用分号、顿号、逗号等方式合并多个条目到同一行。
            - 示例正确格式：
            "1.第一条内容\n2.第二条内容\n3.第三条内容"

            表格单元格理解规则（重要）：
            - 知识片段可能来自表格，每行使用 " | " 分隔单元格。
            - "<空>" 表示空单元格，对应值为 ""。
            - 字段名后不一定是值：
            - 若字段名后是 "<空>" → 值为 ""。
            - 若字段名后是另一个字段名 → 继续向后寻找第一个“非字段名 / 非占位符”的单元格作为值。
            - 示例："字段A | 字段B | 值" → 字段A="", 字段B="值"。

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

    def _build_user_prompt(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}

            【知识库片段】
            {content_text}

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

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

    def node(state: TaskState, config: RunnableConfig):
        task = state["task"]
        path = task.path
        keys = list(task.keys)

        task_id = f"{'/'.join(path) or '<root>'}:{','.join(keys)}"

        default_kv = {k: "" for k in keys}
        default_patch = make_patch(path, default_kv)

        try:
            with limiter:
                # 从 config 中获取 collection_name（运行时参数）
                config_dict = config.get("configurable", {}) if isinstance(config, dict) else {}
                collection_name = config_dict.get("collection_name", "file_watcher_collection")
                
                # 1) 检索 query：路径信息 + keys
                # 把上层标题带进去提升命中
                query_parts = list(path) + keys
                query = "；".join([p for p in query_parts if p])

                retrieve_component = RetrieveComponent(client=client)
                docs = retrieve_component(query=query, collection_name=collection_name, limit=TOP_K)
                docs = slim_docs(docs, MAX_DOC_CHARS)

                # 没 docs：直接返回默认
                if not docs:
                    return {
                        "context_by_task": {task_id: []},
                        "final_result": default_patch,
                    }

                formatted_context = "\n\n".join(
                    [f"片段 {i+1}: {doc}" for i, doc in enumerate(docs)]
                )

                # 2) 一次性抽取 keys（严格 JSON object）
                json_skeleton = "{\n" + ",\n".join([f'  "{k}": ""' for k in keys]) + "\n}"

                system_prompt = _build_system_prompt(keys)
                user_prompt = _build_user_prompt(formatted_context, path, keys)

                llm = LLMComponent(
                    model_id=MODEL_ID,
                    system_prompt=system_prompt,
                    prompt=user_prompt,
                    response_format=ResponseFormat.JSON,
                    temperature=0.0,
                )

                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)

                return {
                    "context_by_task": {task_id: docs},
                    "final_result": patch,
                }

        except Exception as e:
            return {
                "context_by_task": {task_id: []},
                "final_result": default_patch,
                "errors": {task_id: f"{type(e).__name__}: {e}"},
            }

    return node


def create_validate_node():
    def validate_node(state: OverallState):
        res = state.get("final_result", {}) or {}
        missing = []

        def get_in(d: Dict[str, Any], path: Tuple[str, ...]) -> Dict[str, Any]:
            cur = d
            for p in path:
                if not isinstance(cur, dict):
                    return {}
                cur = cur.get(p, {})
            return cur if isinstance(cur, dict) else {}

        for t in state["tasks"]:
            scope = get_in(res, t.path)
            for k in t.keys:
                if not str(scope.get(k, "")).strip():
                    missing.append((".".join(t.path + (k,))) if t.path else k)

        if missing:
            return {"errors": {"missing": "；".join(missing)}}
        return {}
    
    return validate_node


def map_tasks(state: OverallState):
    return [Send("extract_task", {"task": t}) for t in state["tasks"]]


def create_graph(init_args: WorkflowInitArgs | Dict[str, Any]):
    if isinstance(init_args, dict):
        init_args = WorkflowInitArgs(**init_args)
    client = HttpAlayaMemClient(init_args.alayamem_url)
    
    g = StateGraph(OverallState, WorkflowContext, input_type=Input, output_type=Output)
    g.add_node("plan", create_plan_node())
    g.add_node("extract_task", create_extract_task_node(client))
    g.add_node("validate", create_validate_node())
    
    g.add_edge(START, "plan")
    g.add_conditional_edges("plan", map_tasks, ["extract_task"])
    g.add_edge("extract_task", "validate")
    g.add_edge("validate", END)

    return g.compile()

