from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, TypedDict, Annotated, Union, TypeAlias, 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


FieldSpec: TypeAlias = Union[str, Dict[str, List["FieldSpec"]]]  # 递归：dict -> list[FieldSpec]

def merge_dicts(a: Dict, b: Dict) -> Dict:
    return {**a, **b}

def deep_merge(a: Dict[str, Any], b: Dict[str, Any]) -> Dict[str, Any]:
    out = dict(a or {})
    for k, v in (b or {}).items():
        if k in out and isinstance(out[k], dict) and isinstance(v, dict):
            out[k] = deep_merge(out[k], v)
        else:
            out[k] = v
    return out


@dataclass(frozen=True)
class GroupTask:
    path: Tuple[str, ...]   # 父路径，如 ("个人信息","联系方式")；根为 ()
    keys: Tuple[str, ...]   # 该路径下需要抽取的叶子字段名


class OverallState(TypedDict):
    fields: List[FieldSpec]  # 输入模板（递归）
    tasks: List[GroupTask]   # 规划出的任务列表

    # 调试信息：每个任务的检索片段
    context_by_task: Annotated[Dict[str, List[str]], merge_dicts]

    # 最终值树：通过 deep_merge reducer 并发合并 patch
    final_result: Annotated[Dict[str, Any], deep_merge]

    errors: Annotated[Dict[str, str], merge_dicts]


class TaskState(TypedDict):
    task: GroupTask



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

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]] = []

    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 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}


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



def make_patch(path: Tuple[str, ...], kv: Dict[str, str]) -> Dict[str, Any]:
    """
    path=("个人信息","联系方式"), kv={"电话":"..","邮箱":".."} =>
    {"个人信息":{"联系方式":{"电话":"..","邮箱":".."}}}
    """
    node: Dict[str, Any] = dict(kv)
    for p in reversed(path):
        node = {p: node}
    return node



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 create_extract_task_node(
    client: HttpAlayaMemClient,
    *,
    max_concurrency: int = 10,
    top_k: int = 5,
    max_doc_chars: int = 400,
):
    limiter = Semaphore(max_concurrency)

    def slim_docs(docs: List[str]) -> List[str]:
        out = []
        for d in docs or []:
            s = str(d)
            if len(s) > max_doc_chars:
                s = s[:max_doc_chars] + "…"
            out.append(s)
        return out

    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)}"

        # 默认 patch：保证结构稳定（缺失也填空）
        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
                # path 越深，越应该把上层标题带进去提升命中
                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)

                # 没 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_name="deepseek-chat",
                    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 validate_node(state: OverallState, config: RunnableConfig):
    # 简单缺失检查：把 tasks 展开期望字段，看看 final_result 是否为空
    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 {}


# -------------------------
# Build graph
# -------------------------
def create_graph(init_args: Dict[str, Any]):
    client = HttpAlayaMemClient(init_args["alayamem_url"])
    g = StateGraph(OverallState)

    g.add_node("plan", plan_node)
    g.add_node("extract_task", create_extract_task_node(client, max_concurrency=10, top_k=3))
    g.add_node("validate", 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()


if __name__ == "__main__":
    app = create_graph({"alayamem_url": "http://10.16.70.46:5555"})

    input_data: OverallState = {
        "fields": [
            {
                "申请人信息": [
                "姓名",
                "性别",
                "出生年月",
                "民族",
                "学位",
                "职称",
                "是否在站博士后",
                "电子邮箱",
                "办公电话",
                "国别或地区",
                "申请人类别",
                "工作单位",
                "主要研究领域"
                ]
            },
            {
                "依托单位信息": [
                "名称",
                "联系人",
                "电子邮箱",
                "电话",
                "网站地址"
                ]
            },
            {
                "合作研究单位信息": [
                "单位名称"
                ]
            },
            {
                "项目基本信息": [
                "项目名称",
                "英文名称",
                "资助类别",
                "亚类说明",
                "附注说明",
                "申请代码",
                "研究期限",
                "研究方向",
                "申请资助经费",
                "研究属性",
                "中文关键词",
                "英文关键词"
                ]
            },
            "中文摘要",
            "英文摘要"
            ],
        "tasks": [],
        "context_by_task": {},
        "final_result": {},
        "errors": {},
    }

    config = {
        "configurable": {
            "collection_name": "file_watcher_collection",
        }
    }
    out = app.invoke(input_data, config=config)
    print("final_result:")
    print(json.dumps(out["final_result"], ensure_ascii=False, indent=2))
    print("\nerrors:")
    print(json.dumps(out["errors"], ensure_ascii=False, indent=2))
